mayoko’s diary

プロコンとかいろいろ。

AIM Tech Round 3 (Div. 1) C. Centroids

問題

codeforces.com

頂点数 n の木 T 上の点 v が T の centorid であるとは, T から v を取り除いて得られる任意の連結成分が, 全て頂点数 n/2 以下であることを言う。

各頂点について, その頂点が以下の操作を一回以下許したとき centroid になるかどうかを判定せよ。
操作:木のある辺を消して, 好きな辺を加える(ただし新しいグラフも連結でなければならない)。

解法

操作を許さない場合にどうなるかを考えます。この場合はかなり単純です。頂点 v から伸びる辺の頂点から構成される部分木が, 全て n/2 以下の連結成分で構成されているか判定すれば良いです。いちいち頂点 v を根にして木dp しているとまにあわないですが, 0 を頂点にした木のみを考えて, 親方向側の連結成分の個数は n - (v の部分木のサイズ) というようにすれば良いです。

次に操作を許す場合を考えます。頂点 v の部分木でサイズが n/2 より大きいものがあった場合, その部分木の一部を切り取って v に直接つなぎかえるのが最適です。ここで部分木の頂点 ch で 2 通りの場合分けをします。

0 を根とした木でも ch が v の子である場合
この場合は, ch 以下の部分木を単純に調べるだけで良いので, euler_tour とか用意しておいてその中で部分木のサイズが n/2 である最大のものを求めます。ch の部分木のサイズ - この部分木のサイズ が n/2 以下なら操作をして centroid になることになります。

0 を根とした木で ch が v の親である場合
この場合がややこしいです。さらに二通りに分けます。

・つなげる木が v の祖先である場合
この場合, v からつなぎなおす頂点を a とすると, v - ch の木は,

  • (a の部分木サイズ) - (v の部分木サイズ) の木
  • n - (a の部分木サイズ) の木

に別れることがわかります。 n - (a の部分木サイズ) が n/2 以下で最大のものを覚えておけばこの中で (a の部分木サイズ) - (v の部分木サイズ) が n/2 以下になるものがあるかは, O(1) でわかります。

・つなげる木が v の祖先でない場合(要するに v の祖先の子孫のうち v の部分木以外のやつ)
木dp するとき dfs から抜け出すタイミングで部分木のサイズのメモしていくと, 頂点 v にたどり着いたときにメモされているのは, v 祖先の子孫の部分木のサイズのみになっています(直接の先祖はまだ v が子にあるので dfs から抜けておらず, サイズがメモされていない)。

よって, 部分木のサイズが n/2 以下なら 部分木のサイズをメモする, というようにやれば同じようにできます。

この手法でやるときは一回 dfs した後もう一回グラフを逆順にして再び dfs しなければいけないことに注意です。

// セグメント木(RMQ 対応)
// update: k 番目の値を a に変更
// query: [l, r) の区間の最大値を求める
template<typename T>
struct ST {
    vector<T> seg;
    int size;
    ST(int n) {
        size = 1;
        while (size < n) size *= 2;
        seg.resize(2*size-1, 0);
    }
    inline T merge(T x, T y) {
        return max(x, y);
    }
    void update(int k, T a) {
        k += size-1;
        seg[k] = a;
        while (k > 0) {
            k = (k-1)/2;
            seg[k] = merge(seg[k*2+1], seg[k*2+2]);
        }
    }
    T query(int a, int b, int k, int l, int r) {
        if (r <= a || b <= l) return 0;
        if (a <= l && r <= b) return seg[k];
        T vl = query(a, b, k*2+1, l, (l+r)/2);
        T vr = query(a, b, k*2+2, (l+r)/2, r);
        return merge(vl, vr);
    }
    T query(int a, int b) {
        return query(a, b, 0, 0, size);
    }
};

const int MAXN = 400400;
vi G[MAXN];

// Euler-Tour
const int MAXSIZE = MAXN;
int BEGIN[MAXSIZE], END[MAXSIZE];
vector<int> euler_tour;
int K;

void createEulerTour(int v, int p) {
    BEGIN[v] = K++;
    euler_tour.push_back(v);
    for (int el : G[v]) {
        if (el == p) continue;
        createEulerTour(el, v);
        euler_tour.push_back(v);
        K++;
    }
    END[v] = K;
}

int subSize[MAXN];
pii maxi[MAXN];
int n;

int getSize(int v, int p) {
    int& ret = subSize[v];
    pii& ma = maxi[v];
    ma.first = 0, ma.second = -1;
    for (int ch : G[v]) if (ch != p) {
        int tmp = getSize(ch, v);
        if (ma.first < tmp) {
            ma.first = tmp;
            ma.second = ch;
        }
        ret += tmp;
    }
    ret++;
    int tmp = n - ret;
    if (ma.first < tmp) {
        ma.first = tmp;
        ma.second = p;
    }
    return ret;
}

int ans[MAXN];
int tsize;

void dfs(int v, int p, ST<int>& seg, int top) {
    //cout << v << " " << top << " " << tsize << endl;
    if (maxi[v].second == p) {
        if (subSize[top]-subSize[v] <= n/2) ans[v] = 1;
        else if (n-subSize[v]-tsize <= n/2) ans[v] = 1;
    } else {
        int ch = maxi[v].second;
        int tmp = seg.query(BEGIN[ch]+1, END[ch]);
        if (maxi[v].first - tmp <= n/2) ans[v] = 1;
    }
    if (n - subSize[v] <= n/2) top = v;
    for (int ch : G[v]) if (ch != p) {
        dfs(ch, v, seg, top);
    }
    if (subSize[v] <= n/2) tsize = max(tsize, subSize[v]);
}

int main() {
    cin.tie(0);
    ios::sync_with_stdio(false);
    cin >> n;
    for (int i = 0; i < n-1; i++) {
        int a, b;
        cin >> a >> b;
        a--; b--;
        G[a].push_back(b);
        G[b].push_back(a);
    }
    // 部分木のサイズを求める
    getSize(0, -1);
    // euler tour して部分木のサイズをメモ
    createEulerTour(0, -1);
    ST<int> seg(euler_tour.size());
    for (int i = 0; i < n; i++) {
        if (subSize[i] > n/2) {
            seg.update(BEGIN[i], 0);
        } else {
            seg.update(BEGIN[i], subSize[i]);
        }
    }
    dfs(0, -1, seg, 0);
    tsize = 0;
    for (int i = 0; i < n; i++)
        reverse(G[i].begin(), G[i].end());
    dfs(0, -1, seg, 0);
    for (int i = 0; i < n; i++) {
        cout << ans[i] ;
        if (i < n-1) cout << " ";
    }
    cout << endl;
    return 0;
}