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

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

EDPC (Educational DP Contest) V - Subtree (2D)

全方位木 DP の練習問題!!  M が素数とは限らないので「左右から累積和」が必要なタイプの全方位木 DP。

問題概要

正の整数  M と、頂点数  N の木  G が与えられる。この木の各頂点を「黒色」または「白色」に塗る。

各頂点  v について、以下の条件をみたす場合の数を  M で割った余りを求めよ。

  • 頂点  v は黒色である
  • 黒色の頂点は連結である

制約

  •  1 \le N \le 10^{5}
  •  2 \le M \le 10^{9}

 

解法:全方位木 DP

この問題に対して、各頂点  v を根とした木 DP をしてしまうと、 O(N^{2}) の計算量となってしまいます。このようなとき、全方位木 DP がしばしば有効です。

ここでは、通常の「根を 1 つ決めたときの木 DP」を全方位木 DP へと機械的に拡張する方法を示します。

まず、「頂点 0 を黒色にする場合」についての問題を解くとき、次のような木 DP ができます。

  • `dp[v]:頂点  v 根とする根付き木において、頂点  v を黒色に塗った上で黒色頂点全体が連結になっているような塗り方の数に、1 を足した値

とします。ここで「1 を足す」とは、「すべての頂点が白色である場合」を加えることを意味しています。このとき、次のような木 DP ができます。

long long rec(int v, int p) {
    long long res = 1;
    for (auto v2 : G[v]) {
        if (v2 == p) continue;
        res = res * rec(v2, v) % M;
    }
    return (res + 1) % M;
}

この再帰関数 rec() を用いて、rec(0, -1) - 1 が答えとなります。

 

木 DP を抽象化

ここで、上記のような木 DP の抽象化を考えます。次の 4 つのものを定めれば抽象化できそうです。


  • Monoiddp の型
    • 今回は long long
  • IDENTITY:下記の関数 MERGE() についての Monoid の単位元
    • 今回は 1
  • MERGE(Monoid a, Monoid b):部分木と部分木の情報を統合する部分を表す関数
    • 今回は a * b % M を返す
  • ADDNODE(int v, Monoid a):部分木に関して統合した情報 a に対して、最後に頂点 v に関する情報を統合する関数
    • 今回は (a + 1) % M を返す

このように抽象化すると、上の木 DP は次のように書き直せます。

Monoid rec(int v, int p) {
    Monoid res = IDENTITY;
    for (auto v2 : G[v]) {
        if (v2 == p) continue;
        res = MERGE(res, rec(v2, v));
    }
    return ADDNODE(v, res);
}

 

全方位木 DP

この抽象化をもとに、全方位木 DP を考えます。通常の木 DP は、根を 1 つ決めた上で、各頂点を根とする部分木 ( N 個あります) に関する問題を、 O(N) の計算量で順次解いていくものでした。

これに対して全方位木 DP は、木の各辺を除去して得られる部分木 ( 2(N-1) 個あります) に関する問題を、全体  O(N) の計算量で順次解いていくものです。

この  2(N-1) 個の部分木に関する問題が解けているとき、もとの木でどの頂点を根とした場合についても、その子頂点を根とする部分木の情報がすべて得られる状態となっています。これをもとにして、「その頂点を根としたときの問題の解」をならし  O(1) の計算量で求めることができます。

全方位木 DP の詳細については、次の記事に詳しく書かれています。

ei1333.hateblo.jp

ここでは、抽象化した全方位木 DP の実装を示します。

// re-rooting
template<class Monoid> struct ReRooting {
    using Graph = vector<vector<int>>;
    using MergeFunc = function<Monoid(Monoid, Monoid)>;
    using AddNodeFunc = function<Monoid(int, Monoid)>;
    
    // core member
    Graph G;
    Monoid IDENTITY;
    MergeFunc MERGE;
    AddNodeFunc ADDNODE;
    
    // inner data
    vector<vector<Monoid>> dp;
    
    // constructor
    ReRooting() {}
    ReRooting(const Graph &g, const Monoid &identity,
              const MergeFunc &merge, const AddNodeFunc &addnode) {
        G = g;
        IDENTITY = identity;
        MERGE = merge;
        ADDNODE = addnode;
        build();
    }
    
    // re-looting dp
    Monoid rec(int v, int p) {
        Monoid res = IDENTITY;
        dp[v].assign(G[v].size(), IDENTITY);
        for (int i = 0; i < G[v].size(); ++i) {
            int v2 = G[v][i];
            if (v2 == p) continue;
            dp[v][i] = rec(v2, v);
            res = MERGE(res, dp[v][i]);
        }
        return ADDNODE(v, res);
    }
    void rerec(int v, int p, Monoid pval) {
        for (int i = 0; i < G[v].size(); ++i) {
            int v2 = G[v][i];
            if (v2 == p) {
                dp[v][i] = pval;
                continue;
            }
        }
        vector<Monoid> left(G[v].size() + 1, IDENTITY);
        vector<Monoid> right(G[v].size() + 1, IDENTITY);
        for (int i = 0; i < G[v].size(); ++i) {
            left[i + 1] = MERGE(left[i], dp[v][i]);
            right[i + 1] = MERGE(right[i], dp[v][(int)G[v].size() - i - 1]);
        }
        for (int i = 0; i < G[v].size(); ++i) {
            int v2 = G[v][i];
            if (v2 == p) continue;
            Monoid pval2 = MERGE(left[i], right[(int)G[v].size() - i - 1]);
            rerec(v2, v, ADDNODE(v, pval2));
        }
    }
    void build() {
        dp.assign(G.size(), vector<Monoid>());
        int root = 0, nullparent = -1;
        rec(root, nullparent);
        rerec(root, nullparent, IDENTITY);
    }
    
    // getter
    Monoid get(int v) {
        Monoid res = IDENTITY;
        for (int i = 0; i < G[v].size(); ++i) {
            res = MERGE(res, dp[v][i]);
        }
        return ADDNODE(v, res);
    }
    
    // dump
    friend constexpr ostream& operator << (ostream &os, const ReRooting<Monoid> &rr) {
        for (int v = 0; v < rr.G.size(); ++v) {
            for (int i = 0; i < rr.G[v].size(); ++i) {
                os << v << " -> " << rr.G[v][i] << ": " << rr.dp[v][i] << endl;
            }
        }
        return os;
    }
};

 

まとめ

今回の問題に戻ると、木 DP を次のように定義して、全方位木 DP に投げればよいです。全体として計算量は  O(N) となります。

// 木 DP で更新する情報の型
using Monoid = long long;

// 木 DP の更新演算に関する単位元
Monoid identity = 1;

// 木 DP の更新演算
auto merge = [&](Monoid a, Monoid b) -> Monoid { return a * b % M; };

// 木 DP において頂点 v の各部分木の更新をしたあとの頂点 v 自身による更新
auto addnode = [&](int v, Monoid a) -> Monoid { return (a + 1) % M; };

// 全方位木 DP の実施
ReRooting<Monoid> rr(G, identity, merge, addnode);

 

コード

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


// re-rooting
template<class Monoid> struct ReRooting {
    using Graph = vector<vector<int>>;
    using MergeFunc = function<Monoid(Monoid, Monoid)>;
    using AddNodeFunc = function<Monoid(int, Monoid)>;
    
    // core member
    Graph G;
    Monoid IDENTITY;
    MergeFunc MERGE;
    AddNodeFunc ADDNODE;
    
    // inner data
    vector<vector<Monoid>> dp;
    
    // constructor
    ReRooting() {}
    ReRooting(const Graph &g, const Monoid &identity,
              const MergeFunc &merge, const AddNodeFunc &addnode) {
        G = g;
        IDENTITY = identity;
        MERGE = merge;
        ADDNODE = addnode;
        build();
    }
    
    // re-looting dp
    Monoid rec(int v, int p) {
        Monoid res = IDENTITY;
        dp[v].assign(G[v].size(), IDENTITY);
        for (int i = 0; i < G[v].size(); ++i) {
            int v2 = G[v][i];
            if (v2 == p) continue;
            dp[v][i] = rec(v2, v);
            res = MERGE(res, dp[v][i]);
        }
        return ADDNODE(v, res);
    }
    void rerec(int v, int p, Monoid pval) {
        for (int i = 0; i < G[v].size(); ++i) {
            int v2 = G[v][i];
            if (v2 == p) {
                dp[v][i] = pval;
                continue;
            }
        }
        vector<Monoid> left(G[v].size() + 1, IDENTITY);
        vector<Monoid> right(G[v].size() + 1, IDENTITY);
        for (int i = 0; i < G[v].size(); ++i) {
            left[i + 1] = MERGE(left[i], dp[v][i]);
            right[i + 1] = MERGE(right[i], dp[v][(int)G[v].size() - i - 1]);
        }
        for (int i = 0; i < G[v].size(); ++i) {
            int v2 = G[v][i];
            if (v2 == p) continue;
            Monoid pval2 = MERGE(left[i], right[(int)G[v].size() - i - 1]);
            rerec(v2, v, ADDNODE(v, pval2));
        }
    }
    void build() {
        dp.assign(G.size(), vector<Monoid>());
        int root = 0, nullparent = -1;
        rec(root, nullparent);
        rerec(root, nullparent, IDENTITY);
    }
    
    // getter
    Monoid get(int v) {
        Monoid res = IDENTITY;
        for (int i = 0; i < G[v].size(); ++i) {
            res = MERGE(res, dp[v][i]);
        }
        return ADDNODE(v, res);
    }
    
    // dump
    friend constexpr ostream& operator << (ostream &os, const ReRooting<Monoid> &rr) {
        for (int v = 0; v < rr.G.size(); ++v) {
            for (int i = 0; i < rr.G[v].size(); ++i) {
                os << v << " -> " << rr.G[v][i] << ": " << rr.dp[v][i] << endl;
            }
        }
        return os;
    }
};


int main() {
    int N, M;
    cin >> N >> M;
    
    using Graph = vector<vector<int>>;
    Graph G(N);
    for (int i = 0; i < N - 1; ++i) {
        int x, y;
        cin >> x >> y;
        --x, --y;
        G[x].push_back(y);
        G[y].push_back(x);
    }
    
    using Monoid = long long;
    Monoid identity = 1;
    auto merge = [&](Monoid a, Monoid b) -> Monoid { return a * b % M; };
    auto addnode = [&](int v, Monoid a) -> Monoid { return (a + 1) % M; };
    ReRooting<Monoid> rr(G, identity, merge, addnode);
    
    //cout << rr << endl;
    
    for (int v = 0; v < N; ++v) {
        cout << (rr.get(v) + M - 1) % M << endl;
    }
}