けんちょんの競プロ精進記録

競プロの精進記録や小ネタを書いていきます

AtCoder AGC 029 E - Wandering TKHS (赤色, 1200 点)

クエリ処理問題 各kに対して 木 クエリ(木上) グラフ データ構造テク：差分更新 累積和 DP 木DP AtCoder AtCoder1200点 AGC-E 更新箇所が線形オーダー数な二次元DP 赤色diff 変化・遷移先が限られる 累積max

またしても、最後の最後がよく詰めきれず... (でもその最後のところの詰めの大変さが、この難易度帯の特徴なんだよね)

問題へのリンク

問題概要

N 頂点のツリーが与えられる。根ノードの番号を 1 とする。各ノード i = 2, 3, \dots, N について、以下のクエリに答えよ:

  • 初期状態を集合 { i } とする
  • 今得られている集合に含まれるノードに隣接するすべてのノードのうち、番号が最も小さいノードを新たに集合に加える
  • この操作を行ったとき、ノード 1 が集合に加えられるまでの操作回数を答えよ

制約

  • 2 \le N \le 2 × 10^{5}

考えたこと

各クエリを十分高速にやらないといけない。ツリーの各ノードに対して求める問題なのだから、

頂点 v に対して、その親ノード p の結果を利用して高速に求める

という感じになっている気がして来る。そのことを意識した上で、ノード v 単体でクエリに答えようとしたらどうなるのかを考察してみる。その過程で、ノード p の結果を使えそうな部分を模索する感じで。

v から 1 までは全部通っていくわけだが、その過程で一番大きいノード m(v) のやつがボトルネックになりそう。 v より下にあるノードのうち、 v から順に下に下がっていったときに、 m(v) より小さいノードのみを使って辿れるノードはすべて集合に入って行くことになる。

というわけで、

  • f(v, c) := v を根とする部分木の中を、v から出発して番号が c より小さいノードのみを使って辿っていったときに辿れるノードの個数

を高速に求めたいという気持ちになる。でもこれ、 O(N^{2}) かかるような気がしてしまう。。。(ここで詰まってた)

考察の足りてなかったところ

この f(v, c) の計算が、実は必要なところだけをやれば O(N) でできるのだが、それは解説を見て「そうか！」となった。...とその前に、上の考察において、

  • v より少し祖先にあるノード u について、u の部分木 (v 方向以外) の頂点はどのくらい取られるのか

という部分についてはちゃんと考えてなかった。実は一般に、

u の部分木 (v 方向を含まず) のうち、集合に加えられるノードは、u より上にあるノードのうち最大のノード番号を m(u) として、m(u) より小さいノードのみを使って辿れるノード

となることがわかる。この個数は、v について考えているときも、p について考えているときも等しいことがわかる。このことに気づくと明快になる。つまり、

v についての答えと、p についての答えの差分は、v の部分木上のノードにしかない

ということが言える。いよいよ、v についての答えと、p についての答えの差分がどうなるかを考察する。

m(p) < v のとき、すなわち p のどの祖先よりも v の方が大きいとき

p についての答えには、v は一切含まれない。よってその差分は

  • f(v, m(v))

となる。

m(p) > v のとき、すなわち p のある祖先 m(p) が存在してそれが v より大きいとき

p についての答えにも v は含まれる。

最初この場合は、p についての答えと v についての答えが一致すると勘違いしてしまった。だが微妙な差分があることがある。

v の部分木について

  • p についての答えに含まれるノード数は f(v, m(p))
  • v についての答えに含まれるノード数は f(v, m(v))

これらは一般には一致しない。差分は、-f(v, m(p)) + f(v, m(v))

f(v, c) の計算

メモ化再帰で単純に計算できるのだが、一見 O(N^{2}) に見えるが実はすべての c についてこれを計算するわけではないから、全体として O(N) になるという話。

ちょっとこれはなかなか今の僕の力量だと気付き難い。。。

ざっくり、c は根からの葉までいたる経路上の累積 max に登場するものしか使わない。だから根から葉に向かってノード番号が単調増加になっているケースが最悪に見えるのだが、そんなことなくて、f(v, c) のツリー DP を具体的に見ると

map<pint,int> ma;
int rec2(int v, int p, int c) {
    if (ma.count(pint(v, c))) return ma[pint(v, c)];
    int res = 1;
    for (auto ch : G[v]) {
        if (ch == p) continue;
        if (ch > c) continue;
        res += rec2(ch, v, c);
    }
    return ma[pint(v, c)] = res;
}

という感じになるわけだが、if (ch > c) continue; によって一瞬で枝刈りされて終わる。なので結果的に O(N) になる。

どんな場合でも O(N) になる気がして来る。根から葉への経路上の累積 max 値が更新されるたびに、それが葉へと向かうときに、次の累積 max に到達した時点で枝刈りされる。なので、結局全体として、ツリー全体を 1 回舐めるような感じになる (枝刈り時に参照される分を含めると最悪 2 回ずつ)。

実際は、m(v, m(p)) と m(v, m(v)) の二種類が見られるので、ツリー全体を最悪でも 3 回舐める程度の計算量になる。よって全体として O(N) でできる。

#include <iostream>
#include <vector>
#include <map>
using namespace std;
using pint = pair<int,int>;

int N;
vector<vector<int> > G;

map<pint,int> ma;
int rec2(int v, int p, int c) {
    if (ma.count(pint(v, c))) return ma[pint(v, c)];
    int res = 1;
    for (auto ch : G[v]) {
        if (ch == p) continue;
        if (ch > c) continue;
        res += rec2(ch, v, c);
    }
    return ma[pint(v, c)] = res;
}

vector<int> res;
void rec(int v, int p, int mp) {
    if (mp > v) {
        res[v] = res[p] - rec2(v, p, mp) + rec2(v, p, max(mp, p));
        for (auto ch : G[v]) {
            if (ch == p) continue;
            rec(ch, v, max(mp, p));
        }
    }
    else {
        res[v] = res[p] + rec2(v, p, max(mp, p));
        for (auto ch : G[v]) {
            if (ch == p) continue;
            rec(ch, v, max(mp, p));
        }
    }
}

int main() {
    cin >> N;
    G.assign(N, vector<int>());
    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);
    }
    res.assign(N, 0);
    for (auto v : G[0]) rec(v, 0, -1);
    for (int v = 1; v < N; ++v) {
        cout << res[v];
        if (v != N-1) cout << " ";
    }
    cout << endl;
}