unionfind/test/potentialized_unionfind_int.aoj3142.test.cpp

• View this file on GitHub
• Last update: 2024-10-02 23:55:49+09:00
• Problem: https://judge.u-aizu.ac.jp/onlinejudge/description.jsp?id=3142

Depends on

• Potentialized UnionFind （重み付き UnionFind） (unionfind/potentialized_unionfind.hpp)

Code

#define PROBLEM "https://judge.u-aizu.ac.jp/onlinejudge/description.jsp?id=3142"
#include "../potentialized_unionfind.hpp"
#include <algorithm>
#include <iostream>
#include <vector>
using namespace std;

int N;
vector<vector<int>> to;
vector<int> A, B;
PotentializedUnionFind<long long> uf;

long long dfs(int now, int prv) {
    long long acc = B[now] - A[now];
    for (auto nxt : to[now]) {
        if (nxt != prv) {
            long long tmp = dfs(nxt, now);
            uf.unite(nxt, now, tmp);
            acc += tmp;
        }
    }
    return acc;
}

int main() {
    cin.tie(nullptr), ios::sync_with_stdio(false);

    cin >> N;
    uf = PotentializedUnionFind<long long>(N);
    to.resize(N);

    for (int e = 0; e < N - 1; e++) {
        int u, v;
        cin >> u >> v;
        u--, v--;
        to[u].push_back(v);
        to[v].push_back(u);
    }
    A.resize(N);
    B.resize(N);
    for (auto &x : A) cin >> x;
    for (auto &x : B) cin >> x;
    dfs(0, -1);

    long long sum = 0, lo = 1LL << 60;
    for (int i = 0; i < N; i++) {
        long long tmp = uf.diff(0, i);
        sum += tmp, lo = lo > tmp ? tmp : lo;
    }
    cout << sum - lo * N << '\n';
}

#line 1 "unionfind/test/potentialized_unionfind_int.aoj3142.test.cpp"
#define PROBLEM "https://judge.u-aizu.ac.jp/onlinejudge/description.jsp?id=3142"
#line 2 "unionfind/potentialized_unionfind.hpp"
#include <numeric>
#include <utility>
#include <vector>

// Potentialized UnionFind (Weighted UnionFind)
template <class S> struct PotentializedUnionFind {
    std::vector<int> par, sz;
    std::vector<S> pot;
    PotentializedUnionFind(int N = 0) : par(N), sz(N, 1), pot(N) {
        std::iota(par.begin(), par.end(), 0);
    }
    int find(int x) {
        if (par[x] != x) {
            int r = find(par[x]);
            pot[x] = pot[x] + pot[par[x]], par[x] = r;
        }
        return par[x];
    }
    bool unite(int s, int t, S rel_diff) {
        // Relate s and t by f[t] = f[s] + rel_diff
        // Return false iff contradiction happens.
        rel_diff = rel_diff + weight(s) + (-weight(t));
        if ((s = find(s)) == (t = find(t))) return rel_diff == 0;
        if (sz[s] < sz[t]) std::swap(s, t), rel_diff = -rel_diff;
        par[t] = s, sz[s] += sz[t], pot[t] = rel_diff;
        return true;
    }
    S weight(int x) { return find(x), pot[x]; }
    S diff(int s, int t) { return weight(t) + (-weight(s)); } // return f[t] - f[s]
    int count(int x) { return sz[find(x)]; }
    bool same(int s, int t) { return find(s) == find(t); }
};
#line 3 "unionfind/test/potentialized_unionfind_int.aoj3142.test.cpp"
#include <algorithm>
#include <iostream>
#line 6 "unionfind/test/potentialized_unionfind_int.aoj3142.test.cpp"
using namespace std;

int N;
vector<vector<int>> to;
vector<int> A, B;
PotentializedUnionFind<long long> uf;

long long dfs(int now, int prv) {
    long long acc = B[now] - A[now];
    for (auto nxt : to[now]) {
        if (nxt != prv) {
            long long tmp = dfs(nxt, now);
            uf.unite(nxt, now, tmp);
            acc += tmp;
        }
    }
    return acc;
}

int main() {
    cin.tie(nullptr), ios::sync_with_stdio(false);

    cin >> N;
    uf = PotentializedUnionFind<long long>(N);
    to.resize(N);

    for (int e = 0; e < N - 1; e++) {
        int u, v;
        cin >> u >> v;
        u--, v--;
        to[u].push_back(v);
        to[v].push_back(u);
    }
    A.resize(N);
    B.resize(N);
    for (auto &x : A) cin >> x;
    for (auto &x : B) cin >> x;
    dfs(0, -1);

    long long sum = 0, lo = 1LL << 60;
    for (int i = 0; i < N; i++) {
        long long tmp = uf.diff(0, i);
        sum += tmp, lo = lo > tmp ? tmp : lo;
    }
    cout << sum - lo * N << '\n';
}

Back to top page