ABC163F path pass i 解説

はじめに

黄色diff下位くらいかと思ってたら橙だったので、自分の考察を整理する意味も込めて書いてみることにしました。

問題

• ABC163 F - path pass i

---

• 木の各頂点に色が塗られている
• 各 について、色 が塗られている頂点を一度以上通るパスを数える

制約

• 

考察

はじめに (当たり前な話)

パスは 個あるので、パスごとに考えていると間に合いません。(当たり前かもしれないけど、一応)

余事象を考えるという話

色 を通るパスを直接数え上げようとすると、ダブりが出たりしてとてつもなく難しそうですが、こういう時は落ち着いて余事象を求めることを考えるといいです。(この考えは割と大事な気がする)

「色 を一度以上通る」の余事象は「色 を一度も通らない」です。これを数えて、パスの総数から引き算すればいいです。

その余事象はどう求めるのか

超当たり前ですが、グラフに色 の頂点が存在しないなら、その中でどんなパスをとっても色 の頂点を一度も通りません。

なので、色 を通らないパスを数えたいなら、色 の各頂点で木を切断して、各連結成分内で個別に求めてやればいいです。

実は、木の頂点数から、パスの数は一意に定まります (二頂点の選び方と一対一対応するので)。 ちなみに頂点数を とすると パスは 個です。

以上から、色 の頂点で切断したときの連結成分の頂点数を知りたくなります。 これは、木を根付き木にして、各部分木の頂点数を求めておくとできます。

たとえばこのような木の赤い頂点について考えます。根は頂点 です。

赤い頂点で木を切断すると、 という連結成分ができます。

頂点 を根とする部分木の頂点数を とすると、 の方の頂点数は 、 の方は として求められます。

実装の際は、色ごとに stack を用意して dfs をすると、きれいに書ける気がします。(実装例を見てね)

実装

オーバーフローしない！(素振り)

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

constexpr int MAX = 200010;

// 関数にしておくと便利
long long subcalc(long long x) { return x * (x + 1) / 2; }

int n;
int c[MAX];
vector<int> G[MAX];

int sz[MAX];
stack<int> last[MAX];
long long ans[MAX];

// 部分木のサイズを求める
void dfs1(int v, int p) {
    sz[v] = 1;
    for(int u : G[v]) {
        if(u != p) {
            dfs1(u, v);
            sz[v] += sz[u];
        }
    }
}

// 答えの計算
void dfs2(int v, int p) {
    if(v == 0) { // 根にアクセスする前に切断しておく
        for(int i = 1; i <= n; i++)
            last[i].push(n);
        dfs2(1, 0);
        for(int i = 1; i <= n; i++)
            ans[i] += subcalc(last[i].top());
    } else {
        last[c[v]].top() -= sz[v];

        for(int u : G[v]) {
            if(u != p) {
                last[c[v]].push(sz[u]);
                dfs2(u, v);
                ans[c[v]] += subcalc(last[c[v]].top());
                last[c[v]].pop();
            }
        }
    }
}

int main() {
    cin >> n;
    for(int i = 1; i <= n; i++) {
        cin >> c[i];
        last[i].push(n);
    }
    for(int i = 0; i < n - 1; i++) {
        int a, b;
        cin >> a >> b;
        G[a].push_back(b);
        G[b].push_back(a);
    }

    dfs1(1, 0);
    dfs2(0, 0);

    for(int i = 1; i <= n; i++)
        cout << subcalc(n) - ans[i] << endl;

    return 0;
}

おわりに

コンテスト中は stack じゃなくて vector を用いたのですが、vector の方が実行時間もメモリもいい結果になりました、何故でしょう？(その辺詳しくないのでよくわからない...)

あと、全然関係ない話なんですが、解説記事で <bits/stdc++.h> を使うのは大丈夫なんでしょうか？