test/AuxiliaryTree.test.cpp

• View this file on GitHub
• Last update: 2025-07-04 15:57:54+09:00
• Problem: https://onlinejudge.u-aizu.ac.jp/problems/0439

Depends on

• AuxiliaryTree.cpp
• LCA.cpp
• SparseTable.cpp

Code

// competitive-verifier: PROBLEM https://onlinejudge.u-aizu.ac.jp/problems/0439
#include "../AuxiliaryTree.cpp"

#define rep(i, n) for (int i = 0; i < n; i++)
template <class T, class U>
inline bool chmin(T& a, U b) {
    if (a > b) {
        a = b;
        return true;
    }
    return false;
}

int main() {
    cin.tie(0);
    ios::sync_with_stdio(false);
    int n;
    cin >> n;
    vector<int> a(n);
    rep(i, n) {
        cin >> a[i];
        a[i]--;
    }
    vector<vector<int>> g(n);
    rep(i, n - 1) {
        int u, v;
        cin >> u >> v;
        u--, v--;
        g[u].push_back(v);
        g[v].push_back(u);
    }
    auxiliary_tree aux_tree(g);
    vector<vector<int>> vec(n);
    rep(i, n) vec[a[i]].push_back(i);
    vector<int> dist(n), idx(n);
    vector<int> ans(n, 1000000000);
    rep(color, n) {
        if (vec[color].empty()) continue;
        auto [root, vs] = aux_tree.query(vec[color]);
        using P = pair<int, int>;
        priority_queue<P, vector<P>, greater<P>> que;
        for (auto v : vs) dist[v] = n, idx[v] = -1;
        for (int s : vec[color]) dist[s] = 0, idx[s] = s, que.push({dist[s], s});
        while (!que.empty()) {
            auto [d, v] = que.top();
            que.pop();
            if (dist[v] != d) continue;
            for (auto u : aux_tree.G[v]) {
                int c = abs(aux_tree.depth[u] - aux_tree.depth[v]);
                if (chmin(dist[u], d + c)) {
                    idx[u] = idx[v];
                    que.push({dist[u], u});
                }
            }
        }
        for (int v : vs) {
            for (int u : aux_tree.G[v]) {
                if (idx[v] == idx[u]) continue;
                chmin(ans[idx[v]], aux_tree.dist(idx[v], idx[u]));
                chmin(ans[idx[u]], aux_tree.dist(idx[v], idx[u]));
            }
        }
        aux_tree.clear(vs);
    }
    rep(i, n) cout << ans[i] << endl;
}

#line 1 "test/AuxiliaryTree.test.cpp"
// competitive-verifier: PROBLEM https://onlinejudge.u-aizu.ac.jp/problems/0439
#line 1 "AuxiliaryTree.cpp"
#include <bits/stdc++.h>
#line 2 "SparseTable.cpp"
using namespace std;
template <class T, T (*op)(T, T), T (*e)()>
struct sparsetable {
    vector<vector<T>> table;
    vector<int> logtable;
    sparsetable() = default;
    sparsetable(vector<T> v) {
        int len = 0;
        while ((1 << len) <= (int)v.size()) len++;
        table.assign(len, vector<T>(1 << len));
        for (int i = 0; i < (int)v.size(); i++) table[0][i] = v[i];
        for (int i = 1; i < len; i++) {
            for (int j = 0; j + (1 << i) <= (1 << len); j++) {
                table[i][j] = op(table[i - 1][j], table[i - 1][j + (1 << (i - 1))]);
            }
        }
        logtable.resize((int)v.size() + 1);
        for (int i = 2; i < (int)logtable.size(); i++) {
            logtable[i] = logtable[(i >> 1)] + 1;
        }
    }
    T query(int l, int r) {
        assert(l <= r);
        if (l == r) return e();
        int len = logtable[r - l];
        return op(table[len][l], table[len][r - (1 << len)]);
    }
};
template <class T, T (*op)(T, T), T (*e)()>
struct disjointsparsetable {
    vector<vector<T>> table;
    vector<int> logtable;
    disjointsparsetable() = default;
    disjointsparsetable(vector<T> v) {
        int len = 0;
        while ((1 << len) <= v.size()) len++;
        table.assign(len, vector<T>(1 << len, e()));
        for (int i = 0; i < (int)v.size(); i++) table[0][i] = v[i];
        for (int i = 1; i < len; i++) {
            int shift = 1 << i;
            for (int j = 0; j < (int)v.size(); j += shift << 1) {
                int t = min(j + shift, (int)v.size());
                table[i][t - 1] = v[t - 1];
                for (int k = t - 2; k >= j; k--) table[i][k] = op(v[k], table[i][k + 1]);
                if (v.size() <= t) break;
                table[i][t] = v[t];
                int r = min(t + shift, (int)v.size());
                for (int k = t + 1; k < r; k++) table[i][k] = op(table[i][k - 1], v[k]);
            }
        }
        logtable.resize(1 << len);
        for (int i = 2; i < logtable.size(); i++) {
            logtable[i] = logtable[(i >> 1)] + 1;
        }
    }
    T query(int l, int r) {
        if (l == r) return e();
        if (l >= --r) return table[0][l];
        int len = logtable[l ^ r];
        return op(table[len][l], table[len][r]);
    };
};
#line 3 "LCA.cpp"
using namespace std;
pair<int, int> lcatree_op(pair<int, int> a, pair<int, int> b) { return min(a, b); }
pair<int, int> lcatree_e() { return {1000000000, -1}; }
struct lca_tree {
    int n, size;
    vector<int> in, ord, depth;
    disjointsparsetable<pair<int, int>, lcatree_op, lcatree_e> st;
    lca_tree(vector<vector<int>> g, int root = 0) : n((int)g.size()), size(log2(n) + 2), in(n), depth(n, n) {
        depth[root] = 0;
        function<void(int, int)> dfs = [&](int v, int p) {
            in[v] = (int)ord.size();
            ord.push_back(v);
            for (int u : g[v]) {
                if (u == p) continue;
                if (depth[u] > depth[v] + 1) {
                    depth[u] = depth[v] + 1;
                    dfs(u, v);
                    ord.push_back(v);
                }
            }
        };
        dfs(root, -1);
        vector<pair<int, int>> vec((int)ord.size());
        for (int i = 0; i < (int)ord.size(); i++) {
            vec[i] = make_pair(depth[ord[i]], ord[i]);
        }
        st = vec;
    }
    int lca(int u, int v) {
        if (in[u] > in[v]) swap(u, v);
        if (u == v) return u;
        return st.query(in[u], in[v]).second;
    }
    int dist(int u, int v) {
        int l = lca(u, v);
        return depth[u] + depth[v] - 2 * depth[l];
    }
};
#line 3 "AuxiliaryTree.cpp"
using namespace std;
struct auxiliary_tree : lca_tree {
    vector<vector<int>> G;
    auxiliary_tree(vector<vector<int>>& g) : lca_tree(g), G(n) {}
    pair<int, vector<int>> query(vector<int> vs, bool decending = false) {
        // decending:親から子の方向のみ辺を貼る
        assert(!vs.empty());
        sort(vs.begin(), vs.end(), [&](int a, int b) { return in[a] < in[b]; });
        int m = vs.size();
        stack<int> st;
        st.push(vs[0]);
        for (int i = 0; i < m - 1; i++) {
            int w = lca(vs[i], vs[i + 1]);
            if (w != vs[i]) {
                int l = st.top();
                st.pop();
                while (!st.empty() && depth[w] < depth[st.top()]) {
                    if (!decending) G[l].push_back(st.top());
                    G[st.top()].push_back(l);
                    l = st.top();
                    st.pop();
                }
                if (st.empty() || st.top() != w) {
                    st.push(w);
                    vs.push_back(w);
                }
                if (!decending) G[l].push_back(w);
                G[w].push_back(l);
            }
            st.push(vs[i + 1]);
        }
        while (st.size() > 1) {
            int x = st.top();
            st.pop();
            if (!decending) G[x].push_back(st.top());
            G[st.top()].push_back(x);
        }
        // {root,vertex_list}
        return make_pair(st.top(), vs);
    }
    void clear(vector<int> vs) {
        for (int v : vs) G[v].clear();
    }
};
#line 3 "test/AuxiliaryTree.test.cpp"

#define rep(i, n) for (int i = 0; i < n; i++)
template <class T, class U>
inline bool chmin(T& a, U b) {
    if (a > b) {
        a = b;
        return true;
    }
    return false;
}

int main() {
    cin.tie(0);
    ios::sync_with_stdio(false);
    int n;
    cin >> n;
    vector<int> a(n);
    rep(i, n) {
        cin >> a[i];
        a[i]--;
    }
    vector<vector<int>> g(n);
    rep(i, n - 1) {
        int u, v;
        cin >> u >> v;
        u--, v--;
        g[u].push_back(v);
        g[v].push_back(u);
    }
    auxiliary_tree aux_tree(g);
    vector<vector<int>> vec(n);
    rep(i, n) vec[a[i]].push_back(i);
    vector<int> dist(n), idx(n);
    vector<int> ans(n, 1000000000);
    rep(color, n) {
        if (vec[color].empty()) continue;
        auto [root, vs] = aux_tree.query(vec[color]);
        using P = pair<int, int>;
        priority_queue<P, vector<P>, greater<P>> que;
        for (auto v : vs) dist[v] = n, idx[v] = -1;
        for (int s : vec[color]) dist[s] = 0, idx[s] = s, que.push({dist[s], s});
        while (!que.empty()) {
            auto [d, v] = que.top();
            que.pop();
            if (dist[v] != d) continue;
            for (auto u : aux_tree.G[v]) {
                int c = abs(aux_tree.depth[u] - aux_tree.depth[v]);
                if (chmin(dist[u], d + c)) {
                    idx[u] = idx[v];
                    que.push({dist[u], u});
                }
            }
        }
        for (int v : vs) {
            for (int u : aux_tree.G[v]) {
                if (idx[v] == idx[u]) continue;
                chmin(ans[idx[v]], aux_tree.dist(idx[v], idx[u]));
                chmin(ans[idx[u]], aux_tree.dist(idx[v], idx[u]));
            }
        }
        aux_tree.clear(vs);
    }
    rep(i, n) cout << ans[i] << endl;
}

Test cases

Env	Name	Status	Elapsed	Memory
g++	00_sample_00.in	AC	5 ms	4 MB
g++	01_small_00.in	AC	4 ms	4 MB
g++	01_small_01.in	AC	4 ms	4 MB
g++	02_corner_00.in	AC	4 ms	4 MB
g++	02_corner_01.in	AC	4 ms	4 MB
g++	02_corner_02.in	AC	4 ms	4 MB
g++	04_rand_00.in	AC	4 ms	4 MB
g++	04_rand_01.in	AC	4 ms	4 MB
g++	04_rand_02.in	AC	4 ms	4 MB
g++	04_rand_03.in	AC	4 ms	4 MB
g++	04_rand_04.in	AC	4 ms	4 MB
g++	04_rand_05.in	AC	4 ms	4 MB
g++	04_rand_06.in	AC	4 ms	4 MB
g++	04_rand_07.in	AC	4 ms	4 MB
g++	05_large_00.in	AC	6 ms	4 MB
g++	05_large_01.in	AC	6 ms	4 MB
g++	05_large_02.in	AC	6 ms	4 MB
g++	05_large_03.in	AC	6 ms	4 MB
g++	06_huge_00.in	AC	22 ms	9 MB
g++	06_huge_01.in	AC	21 ms	9 MB
g++	07_maximum_00.in	AC	203 ms	59 MB
g++	07_maximum_01.in	AC	192 ms	59 MB
g++	08_extreme_00.in	AC	436 ms	118 MB
g++	08_extreme_01.in	AC	395 ms	118 MB
g++	08_extreme_02.in	AC	380 ms	119 MB
g++	10_long_01.in	AC	372 ms	126 MB
g++	10_long_02.in	AC	373 ms	126 MB
g++	10_long_03.in	AC	376 ms	126 MB
g++	11_long_01.in	AC	359 ms	123 MB
g++	12_long_01.in	AC	378 ms	128 MB
g++	20_star_01.in	AC	355 ms	128 MB
g++	20_star_02.in	AC	358 ms	123 MB
g++	20_star_03.in	AC	348 ms	119 MB
g++	20_star_04.in	AC	348 ms	116 MB
g++	20_star_05.in	AC	359 ms	116 MB
g++	21_star_01.in	AC	374 ms	123 MB
g++	21_star_02.in	AC	360 ms	120 MB
g++	21_star_03.in	AC	375 ms	118 MB
g++	21_star_04.in	AC	395 ms	117 MB
g++	21_star_05.in	AC	432 ms	117 MB
clang++	00_sample_00.in	AC	5 ms	4 MB
clang++	01_small_00.in	AC	4 ms	4 MB
clang++	01_small_01.in	AC	4 ms	4 MB
clang++	02_corner_00.in	AC	4 ms	4 MB
clang++	02_corner_01.in	AC	4 ms	4 MB
clang++	02_corner_02.in	AC	4 ms	4 MB
clang++	04_rand_00.in	AC	4 ms	4 MB
clang++	04_rand_01.in	AC	4 ms	4 MB
clang++	04_rand_02.in	AC	4 ms	4 MB
clang++	04_rand_03.in	AC	4 ms	4 MB
clang++	04_rand_04.in	AC	4 ms	4 MB
clang++	04_rand_05.in	AC	4 ms	4 MB
clang++	04_rand_06.in	AC	4 ms	4 MB
clang++	04_rand_07.in	AC	4 ms	4 MB
clang++	05_large_00.in	AC	6 ms	4 MB
clang++	05_large_01.in	AC	6 ms	4 MB
clang++	05_large_02.in	AC	6 ms	4 MB
clang++	05_large_03.in	AC	6 ms	4 MB
clang++	06_huge_00.in	AC	22 ms	9 MB
clang++	06_huge_01.in	AC	21 ms	9 MB
clang++	07_maximum_00.in	AC	202 ms	59 MB
clang++	07_maximum_01.in	AC	192 ms	59 MB
clang++	08_extreme_00.in	AC	431 ms	118 MB
clang++	08_extreme_01.in	AC	394 ms	118 MB
clang++	08_extreme_02.in	AC	372 ms	119 MB
clang++	10_long_01.in	AC	369 ms	124 MB
clang++	10_long_02.in	AC	377 ms	125 MB
clang++	10_long_03.in	AC	363 ms	125 MB
clang++	11_long_01.in	AC	354 ms	122 MB
clang++	12_long_01.in	AC	367 ms	127 MB
clang++	20_star_01.in	AC	355 ms	126 MB
clang++	20_star_02.in	AC	349 ms	122 MB
clang++	20_star_03.in	AC	348 ms	118 MB
clang++	20_star_04.in	AC	336 ms	116 MB
clang++	20_star_05.in	AC	348 ms	116 MB
clang++	21_star_01.in	AC	352 ms	122 MB
clang++	21_star_02.in	AC	353 ms	120 MB
clang++	21_star_03.in	AC	375 ms	119 MB
clang++	21_star_04.in	AC	384 ms	117 MB
clang++	21_star_05.in	AC	408 ms	117 MB

Back to top page