結果

問題 No.2634 Tree Distance 3
ユーザー Aeren
提出日時 2024-02-17 01:46:18
言語 C++23
(gcc 13.3.0 + boost 1.87.0)
結果
AC  
実行時間 779 ms / 3,000 ms
コード長 10,824 bytes
コンパイル時間 4,391 ms
コンパイル使用メモリ 291,856 KB
実行使用メモリ 81,676 KB
最終ジャッジ日時 2024-09-28 23:26:03
合計ジャッジ時間 37,492 ms
ジャッジサーバーID
（参考情報） 		judge2 / judge3
このコードへのチャレンジ
（要ログイン）
ファイルパターン 結果
sample AC * 2
other AC * 69
権限があれば一括ダウンロードができます

ソースコード

diff #
プレゼンテーションモードにする

// #pragma GCC optimize("O3,unroll-loops")
#include <bits/stdc++.h>
// #include <x86intrin.h>
using namespace std;
#if __cplusplus >= 202002L
using namespace numbers;
#endif
template<class T>
struct graph{
    using Weight_t = T;
    struct Edge_t{
        int from, to;
        T cost;
    };
    int n;
    vector<Edge_t> edge;
    vector<vector<int>> adj;
    function<bool(int)> ignore;
    graph(int n = 1): n(n), adj(n){
        assert(n >= 1);
    }
    graph(const vector<vector<int>> &adj, bool undirected = true): n((int)adj.size()), adj(n){
        assert(n >= 1);
        if(undirected){
            for(auto u = 0; u < n; ++ u) for(auto v: adj[u]) if(u < v) link(u, v);
        }
        else for(auto u = 0; u < n; ++ u) for(auto v: adj[u]) orient(u, v);
    }
    graph(const vector<vector<pair<int, T>>> &adj, bool undirected = true): n((int)adj.size()), adj(n){
        assert(n >= 1);
        if(undirected){
            for(auto u = 0; u < n; ++ u) for(auto [v, w]: adj[u]) if(u < v) link(u, v, w);
        }
        else for(auto u = 0; u < n; ++ u) for(auto [v, w]: adj[u]) orient(u, v, w);
    }
    graph(int n, vector<array<int, 2>> &edge, bool undirected = true): n(n), adj(n){
        assert(n >= 1);
        for(auto [u, v]: edge) undirected ? link(u, v) : orient(u, v);
    }
    graph(int n, vector<tuple<int, int, T>> &edge, bool undirected = true): n(n), adj(n){
        assert(n >= 1);
        for(auto [u, v, w]: edge) undirected ? link(u, v, w) : orient(u, v, w);
    }
    int operator()(int u, int id) const{
        #ifdef LOCAL
        assert(0 <= id && id < (int)edge.size());
        assert(edge[id].from == u || edge[id].to == u);
        #endif
        return u ^ edge[id].from ^ edge[id].to;
    }
    int link(int u, int v, T w = {}){ // insert an undirected edge
        int id = (int)edge.size();
        adj[u].push_back(id), adj[v].push_back(id), edge.push_back({u, v, w});
        return id;
    }
    int orient(int u, int v, T w = {}){ // insert a directed edge
        int id = (int)edge.size();
        adj[u].push_back(id), edge.push_back({u, v, w});
        return id;
    }
    void clear(){
        for(auto [u, v, w]: edge){
            adj[u].clear();
            adj[v].clear();
        }
        edge.clear();
        ignore = {};
    }
    graph transposed() const{ // the transpose of the directed graph
        graph res(n);
        for(auto &e: edge) res.orient(e.to, e.from, e.cost);
        res.ignore = ignore;
        return res;
    }
    int degree(int u) const{ // the degree (outdegree if directed) of u (without the ignoration rule)
        return (int)adj[u].size();
    }
    // The adjacency list is sorted for each vertex.
    vector<vector<int>> get_adjacency_list() const{
        vector<vector<int>> res(n);
        for(auto u = 0; u < n; ++ u) for(auto id: adj[u]){
            if(ignore && ignore(id)) continue;
            res[(*this)(u, id)].push_back(u);
        }
        return res;
    }
    void set_ignoration_rule(const function<bool(int)> &f){
        ignore = f;
    }
    void reset_ignoration_rule(){
        ignore = nullptr;
    }
    friend ostream &operator<<(ostream &out, const graph &g){
        for(auto id = 0; id < (int)g.edge.size(); ++ id){
            if(g.ignore && g.ignore(id)) continue;
            auto &e = g.edge[id];
            out << "{" << e.from << ", " << e.to << ", " << e.cost << "}\n";
        }
        return out;
    }
};
// Requires graph
template<class T>
struct dfs_forest{
    int n;
    T base_dist;
    vector<T> dist;
    vector<int> pv;
    vector<int> pe;
    vector<int> order;
    vector<int> pos;
    vector<int> end;
    vector<int> size;
    vector<int> root_of;
    vector<int> root;
    vector<int> depth;
    vector<int> min_depth;
    vector<int> min_depth_origin;
    vector<int> min_depth_spanning_edge;
    // extra_edge[u]: adjacent edges of u which is not part of the spanning forest found during the last dfs-like run.
    vector<vector<int>> extra_edge;
    vector<int> was;
    dfs_forest(T base_dist = 0): base_dist(base_dist){ }
    void init(int n){
        this->n = n;
        pv.assign(n, -1);
        pe.assign(n, -1);
        order.clear();
        pos.assign(n, -1);
        end.assign(n, -1);
        size.assign(n, 0);
        root_of.assign(n, -1);
        root.clear();
        depth.assign(n, -1);
        min_depth.assign(n, -1);
        min_depth_origin.assign(n, -1);
        min_depth_spanning_edge.assign(n, -1);
        dist.assign(n, base_dist);
        extra_edge.assign(n, {});
        was.assign(n, -2);
        attempt = -1;
    }
    int attempt;
    // O(# of nodes reachable from u)
    template<class U, class F = plus<>>
    void _dfs(const graph<U> &g, int u, F UT = plus<>()){
        depth[u] = 0;
        dist[u] = base_dist;
        root_of[u] = u;
        root.push_back(u);
        pv[u] = pe[u] = -1;
        auto recurse = [&](auto self, int u)->void{
            was[u] = attempt;
            pos[u] = (int)order.size();
            order.push_back(u);
            size[u] = 1;
            min_depth[u] = depth[u];
            min_depth_origin[u] = u;
            min_depth_spanning_edge[u] = -1;
            for(auto id: g.adj[u]){
                if(id == pe[u] || g.ignore && g.ignore(id)) continue;
                int v = g(u, id);
                if(was[v] == attempt){
                    if(min_depth[u] > depth[v]){
                        min_depth[u] = depth[v];
                        min_depth_spanning_edge[u] = id;
                    }
                    if(pe[u] != id) extra_edge[u].push_back(id);
                    continue;
                }
                depth[v] = depth[u] + 1;
                dist[v] = UT(g.edge[id].cost, dist[u]);
                pv[v] = u;
                pe[v] = id;
                root_of[v] = root_of[u];
                self(self, v);
                size[u] += size[v];
                if(min_depth[u] > min_depth[v]){
                    min_depth[u] = min_depth[v];
                    min_depth_origin[u] = min_depth_origin[v];
                    min_depth_spanning_edge[u] = min_depth_spanning_edge[v];
                }
            }
            end[u] = (int)order.size();
        };
        recurse(recurse, u);
    }
    // O(# of nodes reachable from src)
    template<class U, class F = plus<>>
    void dfs(const graph<U> &g, const vector<int> &src, F UT = plus<>()){
        assert(g.n <= n);
        root.clear(), order.clear();
        ++ attempt;
        for(auto u: src){
            assert(0 <= u && u < g.n);
            if(was[u] != attempt) _dfs(g, u, UT);
        }
    }
    // O(n + m)
    template<class U, class F = plus<>>
    void dfs_all(const graph<U> &g, F UT = plus<>()){
        assert(g.n <= n);
        root.clear(), order.clear();
        ++ attempt;
        for(auto u = 0; u < g.n; ++ u) if(was[u] != attempt) _dfs(g, u, UT);
    }
    // Check if u is visited during the last dfs-like call.
    bool visited(int u) const{
        assert(0 <= u && u < n);
        return was[u] == attempt;
    }
    // Check if u is an ancestor of v in some spanning tree.
    bool ancestor_of(int u, int v) const{
        assert(visited(u) && visited(v));
        return pos[u] <= pos[v] && end[v] <= end[u];
    }
    // Check if any cycle is found during the last dfs-like call.
    // If must_contain_root = true, the sought cycle is forced to contain one of the root of the trees.
    template<class U>
    optional<pair<int, vector<int>>> find_any_cycle(const graph<U> &g, bool must_contain_root = false) const{
        for(auto u: order) for(auto id: extra_edge[u]){
            int v = g(u, id);
            if(!ancestor_of(v, u) || must_contain_root && root_of[v] != v) continue;
            vector<int> cycle_edge;
            while(u != v) cycle_edge.push_back(pe[u]), u = pv[u];
            reverse(cycle_edge.begin(), cycle_edge.end());
            cycle_edge.push_back(id);
            return pair{v, cycle_edge};
        }
        return {};
    }
};
template<int h = 20>
struct binary_lifting{
    int n = 0;
    vector<int> depth;
    vector<array<int, h>> lift;
    binary_lifting(){ }
    // pv: parent vertex (-1 if root of an arborescence)
    binary_lifting(const vector<int> &pv): n((int)pv.size()), depth(n, numeric_limits<int>::max()), lift(n){
        for(auto u = 0; u < n; ++ u) lift[u][0] = ~pv[u] ? pv[u] : u;
        for(auto bit = 1; bit < h; ++ bit) for(auto u = 0; u < n; ++ u) lift[u][bit] = lift[lift[u][bit - 1]][bit - 1];
    }
    // Requires graph
    template<class Graph>
    binary_lifting(const Graph &g, const vector<int> &roots){
        vector<int> pv(g.n, -1), depth(g.n);
        auto dfs = [&](auto self, int u, int pe)->void{
            for(auto id: g.adj[u]){
                if(id == pe || g.ignore && g.ignore(id)) continue;
                auto &e = g.edge[id];
                int v = u ^ e.from ^ e.to;
                depth[v] = depth[u] + 1;
                pv[v] = u;
                self(self, v, id);
            }
        };
        for(auto u: roots) assert(!depth[u]), pv[u] = u, dfs(dfs, u, -1);
        *this = binary_lifting(pv, depth);
    }
    // pv: parent vertex (-1 if root of an arborescence)
    binary_lifting(const vector<int> &pv, const vector<int> &depth): n((int)pv.size()), depth(depth){
        lift.resize(n);
        for(auto u = 0; u < n; ++ u) lift[u][0] = ~pv[u] ? pv[u] : u;
        for(auto d = 1; d < h; ++ d) for(auto u = 0; u < n; ++ u) lift[u][d] = lift[lift[u][d - 1]][d - 1];
    }
    // Index becomes the current number of nodes
    // O(log n)
    int add_root(){
        int u = n ++;
        depth.push_back(0);
        lift.emplace_back();
        fill(lift.back().begin(), lift.back().end(), u);
        return u;
    }
    // Index becomes the current number of nodes
    // O(log n)
    int add_child(int p){
        assert(0 <= p && p < n);
        int u = n ++;
        depth.push_back(depth[p] + 1);
        lift.emplace_back();
        lift[u][0] = p;
        for(auto d = 1; d < h; ++ d) lift[u][d] = lift[lift[u][d - 1]][d - 1];
    }
    // Get the k-th ancestor of u
    // O(log n)
    int ancestor(int u, int k) const{
        for(auto d = 0; d < h; ++ d) if(k & 1 << d) u = lift[u][d];
        return u;
    }
    // Assumes u and v lies on the same arboresence
    // O(log n)
    int lca(int u, int v) const{
        if(depth[u] < depth[v]) swap(u, v);
        u = ancestor(u, depth[u] - depth[v]);
        if(u == v) return u;
        for(auto d = h - 1; d >= 0; -- d) if(lift[u][d] != lift[v][d]) u = lift[u][d], v = lift[v][d];
        return lift[u][0];
    }
    // Get # of edges between u and v
    // Assumes u and v lies on the same arboresence
    // O(log n)
    int steps(int u, int v, int w = -1) const{
        return depth[u] + depth[v] - 2 * depth[~w ? w : lca(u, v)];
    }
    // For an ancestor p of u, pred(p) is T, ..., T, F, ..., F in decreasing order of depth. Returns the highest p with T
    // O(log n)
    int find_highest(int u, auto pred) const{
        assert(pred(u));
        for(auto d = h - 1; d >= 0; -- d) if(pred(lift[u][d])) u = lift[u][d];
        return u;
    }
};
int main(){
    cin.tie(0)->sync_with_stdio(0);
    cin.exceptions(ios::badbit | ios::failbit);
    int n;
    cin >> n;
    map<int, vector<int>> mp;
    for(auto u = 0; u < n; ++ u){
        int x;
        cin >> x;
        mp[x].push_back(u);
    }
    graph<int> g(n);
    for(auto i = 0; i < n - 1; ++ i){
        int u, v;
        cin >> u >> v, -- u, -- v;
        g.link(u, v, 1);
    }
    binary_lifting lift(g, {0});
    dfs_forest<int> df;
    df.init(n);
    df.dfs(g, {0});
    int x = -1, y = -1;
    vector<int> active(n);
    auto color_path = [&](int u, int v)->void{
        assert(!active[u] && active[v]);
        while(!active[u] && !df.ancestor_of(u, v)){
            active[u] = true;
            u = df.pv[u];
        }
        if(active[u]){
            return;
        }
        v = lift.find_highest(v, [&](int v){ return active[v]; });
        while(v != u){
            v = df.pv[v];
            active[v] = true;
        }
    };
    vector<int> res(n, -1);
    for(auto [_, a]: mp | ranges::views::reverse){
        int i = 0;
        if(!~x){
            x = y = a[i ++];
            active[x] = true;
        }
        for(; i < (int)a.size(); ++ i){
            int u = a[i];
            if(active[u]){
                continue;
            }
            color_path(u, x);
            if(lift.steps(x, y) < lift.steps(x, u)){
                swap(y, u);
            }
            if(lift.steps(x, y) < lift.steps(y, u)){
                swap(x, u);
            }
        }
        for(auto u: a){
            res[u] = max(lift.steps(x, u), lift.steps(y, u));
        }
    }
    ranges::copy(res, ostream_iterator<int>(cout, " "));
    cout << "\n";
    return 0;
}
/*
*/
הההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההה
XXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXX
0