verify/aux_tree.test.cpp

• View this file on GitHub
• Last update: 2025-02-20 17:16:43+09:00
• Problem: https://judge.u-aizu.ac.jp/onlinejudge/description.jsp?id=0439

Depends on

• HLD (Algorithm/hld.hpp)
• verify用テンプレート (Utility/template.hpp)

Code

#define PROBLEM "https://judge.u-aizu.ac.jp/onlinejudge/description.jsp?id=0439"
#include "../Utility/template.hpp"
#include "../Algorithm/hld.hpp"

const int inf = 1001001001;
int main() {
    int n;
    cin >> n;
    vec<int> C(n);
    rep(i, 0, n) cin >> C[i], C[i]--;
    vec<vec<int>> G(n);
    rep(i, 0, n-1) {
        int u, v;
        cin >> u >> v;
        u--, v--;
        G[u].push_back(v);
        G[v].push_back(u);
    }

    vec<vec<int>> dp(n);
    vec<int> ans(n, inf);
    HLD hld(G, 0);

    vec<vec<int>> cv(n);
    rep(i, 0, n) cv[C[i]].push_back(i);

    vec<vec<int>> g(n);

    rep(c, 0, n) {
        int r = hld.aux_tree(cv[c], g);
        if(r == -1) continue;

        auto dfs = [&](auto f, int v, int p) -> int {
            int res = inf;
            int d = g[v].size();
            dp[v].resize(d);
            
            rep(i, 0, d) {
                int to = g[v][i];
                if(to == p) continue;
                dp[v][i] = f(f, to, v);
                int nw = dp[v][i] + hld.dist(v, to);//ここで到達済み
                chmin(res, nw);
            }
            //res++;
            if(C[v] == c) res = 0;
            return res;
        };

        auto bfs = [&](auto f, int v,  int par, int p) -> void {
            int d =  g[v].size();
            rep(i, 0, d) if(g[v][i] == p) dp[v][i] = par;

            vec<int> L(d+1, inf);
            vec<int> R(d+1, inf);
            rep(i, 0, d) {
                int nw = dp[v][i] + hld.dist(v, g[v][i]);
                L[i+1] = min(L[i], nw);
            }

            rrep(i, 0, d) {
                int nw = dp[v][i] + hld.dist(v, g[v][i]);
                R[i] = min(R[i+1], nw);
            }

            if(C[v] == c) {
                ans[v] = L[d];
            }

            rep(i, 0, d) {
                int to = g[v][i];
                if(to == p) continue;
                int nx = min(L[i], R[i+1]);
                if(C[v] == c) nx = 0;
                f(f, to, nx, v);
            }
        };

        dfs(dfs, r, -1);
        bfs(bfs, r, inf, -1);
        hld.clean(g);
    }

    rep(i, 0, n) cout << ans[i] << '\n';
        

}

#line 1 "verify/aux_tree.test.cpp"
#define PROBLEM "https://judge.u-aizu.ac.jp/onlinejudge/description.jsp?id=0439"
#line 1 "Utility/template.hpp"
#include <bits/stdc++.h>
using namespace std;
using ll = long long;
#define rep(i, s, t) for (ll i = s; i < (ll)(t); i++)
#define rrep(i, s, t) for (ll i = (ll)(t) - 1; i >= (ll)(s); i--)
#define all(x) begin(x), end(x)

#define TT template <typename T>
TT using vec = vector<T>;
template <class T1, class T2> bool chmin(T1 &x, T2 y) {
    return x > y ? (x = y, true) : false;
}
template <class T1, class T2> bool chmax(T1 &x, T2 y) {
    return x < y ? (x = y, true) : false;
}
struct io_setup {
    io_setup() {
        ios::sync_with_stdio(false);
        std::cin.tie(nullptr);
        cout << fixed << setprecision(15);
    }
} io_setup;

/*
@brief verify用テンプレート
*/
#line 1 "Algorithm/hld.hpp"
struct HLD {
    using vi = vec<int>;
    using pi = pair<int, int>;
    using pll = pair<long long, long long>;
    vi in, out, par, root, rev, dep, pre_vs;
    vec<ll> dep_w;
    //          親/成分のtop/inの中身→頂点番号
    int n, r;  // 頂点数、根
    
    static vec<vec<int>> extract_graph(const vec<vec<pll>> &G) {
        vec<vec<int>> g(G.size());
        for (int i = 0; i < int(G.size()); i++) {
            for (auto [w, to] : G[i])
                if (i < to) {
                    g[i].push_back(to);
                    g[to].push_back(i);
                }
        }
        return g;
    }
    HLD(const vec<vec<pll>> &g, int a) : HLD(extract_graph(g), a) {
        auto dfs = [&](auto f, int v) -> void {
            for (auto [w, to] : g[v])
                if (to != par[v]) {
                    dep_w[to] = dep_w[v] + w;
                    f(f, to);
                }
        };
        dfs(dfs, r);
    }

    HLD(vec<vi> g, int a) : n(g.size()), r(a) {
        vi siz(n, 0);
        in = out = root = rev = vi(n);
        par = vi(n, -1);
        dep = vi(n, 0);
        dep_w = vec<ll>(n, 0);
        root[r] = r;

        auto dfs_siz = [&](auto f, int v) -> void {
            siz[v]++;
            for (int &to : g[v])
                if (to != par[v]) {
                    dep[to] = dep[v] + 1;
                    par[to] = v;
                    f(f, to);
                    siz[v] += siz[to];
                    if (siz[to] > siz[g[v][0]] || g[v][0] == par[v])
                        swap(to, g[v][0]);
                }
            return;
        };

        dfs_siz(dfs_siz, r);

        int t = 0;

        auto dfs_hld = [&](auto f, int v) -> void {
            rev[t] = v;
            in[v] = t++;
            for (int to : g[v])
                if (to != par[v]) {
                    root[to] = (to == g[v][0] ? root[v] : to);
                    f(f, to);
                }
            out[v] = t;
        };

        dfs_hld(dfs_hld, r);
        for (int i = 0; i < n; i++) dep_w[i] = dep[i];
    }

    // 以下、欲しいもののみ書く

    int operator()(int v) const { return in[v]; }
    int operator()(int u, int v) const {
        assert(par[u] == v || par[v] == u);
        if(par[u] == v) return in[u];
        else return in[v];
    }

    int lca(int a, int b) {
        while (1) {
            if (in[a] > in[b]) swap(a, b);
            if (root[a] == root[b]) return a;
            b = par[root[b]];
        }
    }

    ll dist(int a, int b) {
        int lc = lca(a, b);
        return dep_w[a] + dep_w[b] - 2 * dep_w[lc];
    }

    vec<pi> path(int s, int t, bool edge) {
        vec<pi> ls, rs;
        while (root[s] != root[t]) {
            if (dep[root[s]] > dep[root[t]]) {
                ls.emplace_back(in[s] + 1, in[root[s]]);  // 上り
                s = par[root[s]];
            } else {
                rs.emplace_back(in[root[t]], in[t] + 1);  // 下り
                t = par[root[t]];
            }
        }

        if (dep[s] > dep[t])
            ls.emplace_back(in[s] + 1, in[t] + edge);  // 上り
        else
            rs.emplace_back(in[s] + edge, in[t] + 1);  // 下り

        reverse(all(rs));
        for (auto &p : rs) ls.push_back(p);
        return ls;
    }

    pi subtree(int u, bool edge) { return pi(in[u] + edge, out[u]); }

    int kth_ancestor(int v, int k) {
        if (k > dep[v]) return -1;
        while (v >= 0) {
            if (k <= dep[v] - dep[root[v]]) {
                return rev[in[v] - k];
            } else {
                k -= dep[v] - dep[root[v]] + 1;
                v = par[root[v]];
            }
        }
    }

    int jump(int s, int t, int k) {
        int m = lca(s, t);
        int le = dep[s] - dep[m];
        int ri = dep[t] - dep[m];
        if (0 <= k && k <= le + ri) {
            if (k < le)
                return kth_ancestor(s, k);
            else
                return kth_ancestor(t, le + ri - k);
        }
        return -1;
    }

    int aux_tree(vi vs, vec<vi> &g) {
        if (vs.empty()) return -1;

        auto cmp = [&](int i, int j) { return in[i] < in[j]; };
        sort(all(vs), cmp);
        int m = vs.size();

        rep(i, 0, m - 1) vs.push_back(lca(vs[i], vs[i + 1]));
        sort(all(vs), cmp);
        vs.erase(unique(all(vs)), vs.end());

        vi st;
        for (auto v : vs) {
            while (st.size()) {
                int p = st.back();
                if (in[p] < in[v] && in[v] < out[p]) break;
                st.pop_back();
            }
            if (st.size()) {
                g[st.back()].push_back(v);
                g[v].push_back(st.back());
            }
            st.push_back(v);
        }

        swap(vs, pre_vs);
        return pre_vs[0];
    }

    void clean(vec<vi> &g) {
        for (auto v : pre_vs) g[v] = vi();
        pre_vs = vi();
        return;
    }
};
/*
@brief HLD
@docs doc/hld.md
*/
#line 4 "verify/aux_tree.test.cpp"

const int inf = 1001001001;
int main() {
    int n;
    cin >> n;
    vec<int> C(n);
    rep(i, 0, n) cin >> C[i], C[i]--;
    vec<vec<int>> G(n);
    rep(i, 0, n-1) {
        int u, v;
        cin >> u >> v;
        u--, v--;
        G[u].push_back(v);
        G[v].push_back(u);
    }

    vec<vec<int>> dp(n);
    vec<int> ans(n, inf);
    HLD hld(G, 0);

    vec<vec<int>> cv(n);
    rep(i, 0, n) cv[C[i]].push_back(i);

    vec<vec<int>> g(n);

    rep(c, 0, n) {
        int r = hld.aux_tree(cv[c], g);
        if(r == -1) continue;

        auto dfs = [&](auto f, int v, int p) -> int {
            int res = inf;
            int d = g[v].size();
            dp[v].resize(d);
            
            rep(i, 0, d) {
                int to = g[v][i];
                if(to == p) continue;
                dp[v][i] = f(f, to, v);
                int nw = dp[v][i] + hld.dist(v, to);//ここで到達済み
                chmin(res, nw);
            }
            //res++;
            if(C[v] == c) res = 0;
            return res;
        };

        auto bfs = [&](auto f, int v,  int par, int p) -> void {
            int d =  g[v].size();
            rep(i, 0, d) if(g[v][i] == p) dp[v][i] = par;

            vec<int> L(d+1, inf);
            vec<int> R(d+1, inf);
            rep(i, 0, d) {
                int nw = dp[v][i] + hld.dist(v, g[v][i]);
                L[i+1] = min(L[i], nw);
            }

            rrep(i, 0, d) {
                int nw = dp[v][i] + hld.dist(v, g[v][i]);
                R[i] = min(R[i+1], nw);
            }

            if(C[v] == c) {
                ans[v] = L[d];
            }

            rep(i, 0, d) {
                int to = g[v][i];
                if(to == p) continue;
                int nx = min(L[i], R[i+1]);
                if(C[v] == c) nx = 0;
                f(f, to, nx, v);
            }
        };

        dfs(dfs, r, -1);
        bfs(bfs, r, inf, -1);
        hld.clean(g);
    }

    rep(i, 0, n) cout << ans[i] << '\n';
        

}

Back to top page