Codeforces Round #292 (Div. 1) D. Drazil and Morning Exercise

問題概要

木 T が与えられる.また,クエリが q 回投げられる.それぞれのクエリは L を投げてくる.
木 T の頂点 i に対して,i からそこから最も遠い葉までの距離を D(i) とする.
各クエリに対し,T の部分木 S で,かつ S の任意の二頂点 i, j が |D(i) - D(j)| <= L となるようなものの中で,最も大きいものの大きさを求めよ.
D(i) はあくまでも T での距離であって,S での距離ではないことに注意.

・制約
1 <= n <= 10^5
1 <= q <= 50
1 <= L <= 10^11

解法

全方位木DP,Union-Find木で解ける.
T においてもっとも遠い葉までの距離を,全方位木DPで求める.
そうして得られた距離 d[i] を,降順にソートする.
あとはソートして得られた距離の列上で,前からしゃくとりする.最大が x なら x - L までは含んで良い.
しゃくとりで新たに頂点を追加する場合,その頂点の隣接点で,今現在追加されている頂点があるなら,そこと union-find でマージする.
しゃくとりで左端を縮めるときは,union-find のサイズを1つ小さくするだけでよい.
これは,d を降順でソートしているからである.もし昇順にしてしまうと,サイズを1減らすだけではおかしな値になる.
(イメージとしては,葉の方から根の方にどんどん集合が集まってくる感じ.根からスタートすると,途中で分岐したときに集合を完全にわけないとだめになって,サイズを 1 減らすだけでは対応できない.葉の方からスタートすると,途中でくっつくことがあっても別れることはない.)

ソースコード

#include <bits/stdc++.h>
using namespace std;

using ll = long long;
using pll = pair<ll, ll>;

struct edge {
    int to;
    ll cost;
};

using edges = vector<edge>;
using graph = vector<edges>;

class union_find {
public:
    union_find(int n)
        : par_(n, -1)
    {}

    int root(int x) {
        return par_[x] < 0 ? x : par_[x] = root(par_[x]);
    }

    bool unite(int x, int y) {
        x = root(x); y = root(y);
        if(x == y) {
            return false;
        } else {
            if(par_[x] < par_[y]) {
                par_[x] += par_[y];
                par_[y] = x;
            } else {
                par_[y] += par_[x];
                par_[x] = y;
            }
            return true;
        }
    }

    bool same(int x, int y) {
        return root(x) == root(y);
    }

    int size(int x) {
        return -par_[root(x)];
    }

    void dec(int x) {
        par_[root(x)] += 1;
    }

private:
    std::vector<int> par_;
};

void dfs(int v, int prev, graph const& g, vector<ll>& d) {
    d[v] = 0;
    for(auto& e : g[v]) {
        if(e.to == prev) {
            continue;
        }
        dfs(e.to, v, g, d);
        d[v] = max(d[v], d[e.to] + e.cost);
    }
}

void dfs2(int v, int prev, ll par_d, graph const& g, vector<ll> const& d, vector<pll>& res) {
    res[v].second = v;
    vector<pll> child = {{0, -1}};
    for(auto& e : g[v]) {
        if(e.to == prev) {
            child.emplace_back(par_d + e.cost, e.to);
        } else {
            child.emplace_back(d[e.to] + e.cost, e.to);
        }
    }
    sort(child.rbegin(), child.rend());
    res[v].first = child[0].first;

    for(auto& e : g[v]) {
        if(e.to == prev) {
            continue;
        }
        if(e.to == child[0].second) {
            dfs2(e.to, v, child[1].first, g, d, res);
        } else {
            dfs2(e.to, v, child[0].first, g, d, res);
        }
    }
}

int main() {
    int n;
    cin >> n;
    graph g(n);
    for(int i = 0; i < n - 1; ++i) {
        int x, y;
        ll v;
        cin >> x >> y >> v;
        x--;
        y--;
        g[x].push_back(edge{y, v});
        g[y].push_back(edge{x, v});
    }
    vector<ll> tmp_d(n);
    vector<pll> d(n);
    dfs(0, -1, g, tmp_d);
    dfs2(0, -1, 0, g, tmp_d, d);
    sort(d.rbegin(), d.rend());

    int q;
    cin >> q;
    while(q--) {
        ll L;
        cin >> L;
        int l = 0, r = 0;
        vector<bool> used(n);
        union_find uf(n);
        int res = 0;
        while(r < n) {
            while(r < n && d[l].first <= d[r].first + L) {
                for(auto& e : g[d[r].second]) {
                    if(used[e.to]) {
                        uf.unite(e.to, d[r].second);
                    }
                }
                used[d[r].second] = true;
                r++;
            }
            res = max(res, uf.size(d[l].second));
            uf.dec(d[l].second);
            used[d[l].second] = false;
            l++;
        }
        cout << res << endl;
    }
}

感想

オーバーフローに気づかなくで大量に時間を溶かしてしまった.
UnionFind ってこういう使い方もあるんだなあと勉強になった.