全方位木 DP の抽象化

全方位木 DP の抽象化を書いてみた。
コードと使用例のみ。(使用例はネタバレあり)

気が向いたら解説を書くかもしれない、未来の自分に期待。
少しだけ書くと、rerooting_function では、子ノードの subtree_dp の累積を両側から取って、それぞれの子ノードに対して rerooting_dp を求めるということをしている。
( 2021/02/16 修正 : E - 限界集落 のように辺に向きがある場合に困ったので、子ノードに根方向の部分木の情報を伝播させて、ノードに訪れたときに rerooting_dp を更新するように変更。また、重み付きの方に有向辺を入れる関数も追加。)
最後の参考文献に上げたリンク先を参照すると分かりやすいかもしれない。
このような実装が不必要な場合もあるが、使用例の問題のように割り算ができない場合にも対応するためにもこの実装にした。

あまり多くの問題を見てないので、抽象化がこれでいいのか分からない。
適宜訂正していく可能性が大いにある。
※頂点重み付きも追加 (2021/04/04)

コード

重みなし

▶ソースコードを展開する

// SubtreeMerge : 部分木をマージする関数 (累積を取るときにも用いる)
// PrefixSuffixMerge : prefix と suffix をマージする関数 (rerooting_dp の値を求めるときに用いる)
// IE : 単位元
// subtree_dp[v] := root が根である木で、v を根とした部分木における dp の値
// rerooting_dp[v] := v を根としたときの dp の値

template <typename T>
struct ReRooting {
    using SubtreeMerge = function<T(T, T)>;
    using PrefixSuffixMerge = function<T(T, T)>;

    vector<vector<int>> g;
    const SubtreeMerge f1;
    const PrefixSuffixMerge f2;
    const T IE;     // Identity Element
    vector<T> subtree_dp, rerooting_dp;
    
    ReRooting(int V, const SubtreeMerge s, const PrefixSuffixMerge p, const T I)
        : g(V), f1(s), f2(p), IE(I), subtree_dp(V, I), rerooting_dp(V, I) {}

    void add_edge(int u, int v){
        g[u].emplace_back(v);
        g[v].emplace_back(u);
    }

    // よくある部分木 DP
    void subtree_dfs(int from, int par){
        for(int to : g[from]){
            if(to != par){
                subtree_dfs(to, from);
                subtree_dp[from] = f1(subtree_dp[from], subtree_dp[to]);
            }
        }
    }

    // 各頂点を根としたときの DP
    // par_val : 根方向の部分木の値
    void rerooting_function(int cur, T par_val, int par){
        int V = g[cur].size();
        vector<T> prefix(V + 1, IE), suffix(V + 1, IE);
        for(int i = 0; i < V; ++i){
            int child = g[cur][i];
            if(child == par)    prefix[i + 1] = f1(prefix[i], par_val);
            else    prefix[i + 1] = f1(prefix[i], subtree_dp[child]);
        }
        for(int i = V - 1; i >= 0; --i){
            int child = g[cur][i];
            if(child == par)    suffix[i] = f1(suffix[i + 1], par_val);
            else    suffix[i] = f1(suffix[i + 1], subtree_dp[child]);
        }
        for(int i = 0; i < V; ++i){
            int child = g[cur][i];
            if(child == par){
                rerooting_dp[cur] = f1(subtree_dp[cur], par_val);
            }
            else{
                rerooting_function(child, f2(prefix[i], suffix[i + 1]), cur);
            }
        }
    }

    vector<T> build(int root = 0){
        subtree_dfs(root, -1);
        rerooting_function(root, IE, -1);
        rerooting_dp[root] = subtree_dp[root];
        return rerooting_dp;
    }
};

重み付き

▶ソースコードを展開する

// Mapping : dp + weight を定義する関数

template <typename Weight, typename T>
struct WeightedReRooting {
    struct edge {
        int v;
        Weight w;
        edge() = default;
        edge(int v, Weight w) : v(v), w(w) {}
    };

    using SubtreeMerge = function<T(T, T)>;
    using PrefixSuffixMerge = function<T(T, T)>;
    using Mapping = function<T(T, Weight)>;

    vector<vector<edge>> g;
    const SubtreeMerge f1;
    const PrefixSuffixMerge f2;
    const Mapping f3;
    const T IE;
    vector<T> subtree_dp, rerooting_dp;
    
    WeightedReRooting(int V, const SubtreeMerge s, const PrefixSuffixMerge p, const Mapping m, const T I)
        : g(V), f1(s), f2(p), f3(m), IE(I), subtree_dp(V, I), rerooting_dp(V, I) {}

    void add_edge(int u, int v, Weight w){
        g[u].emplace_back(v, w);
        g[v].emplace_back(u, w);
    }

    void add_directed_edge(int u, int v, Weight w){
        g[u].emplace_back(v, w);
    }

    void subtree_dfs(int from, int par){
        for(auto& [to, weight] : g[from]){
            if(to != par){
                subtree_dfs(to, from);
                subtree_dp[from] = f1(subtree_dp[from], f3(subtree_dp[to], weight));
            }
        }
    }

    void rerooting_function(int cur, T par_val, int par){
        int V = g[cur].size();
        vector<T> prefix(V + 1, IE), suffix(V + 1, IE);
        for(int i = 0; i < V; ++i){
            auto& [child, weight] = g[cur][i];
            if(child == par)    prefix[i + 1] = f1(prefix[i], f3(par_val, weight));
            else    prefix[i + 1] = f1(prefix[i], f3(subtree_dp[child], weight));
        }
        for(int i = V - 1; i >= 0; --i){
            auto& [child, weight] = g[cur][i];
            if(child == par)    suffix[i] = f1(suffix[i + 1], f3(par_val, weight));
            else    suffix[i] = f1(suffix[i + 1], f3(subtree_dp[child], weight));
        }
        for(int i = 0; i < V; ++i){
            auto& [child, weight] = g[cur][i];
            if(child == par){
                rerooting_dp[cur] = f1(subtree_dp[cur], f3(par_val, weight));
            }
            else{
                rerooting_function(child, f2(prefix[i], suffix[i + 1]), cur);
            }
        }
    }

    vector<T> build(int root = 0){
        subtree_dfs(root, -1);
        rerooting_function(root, IE, -1);
        rerooting_dp[root] = subtree_dp[root];
        return rerooting_dp;
    }
};

頂点重み付き

▶ソースコードを展開する

// set(v, w) : 頂点 v の重み w を初期化

template <typename T, typename Weight>
struct VertexWeightedReRooting {
    using SubtreeMerge = function<T(T, T)>;
    using PrefixSuffixMerge = function<T(T, T)>;
    using Mapping = function<T(T, Weight)>;

    vector<vector<int>> g;
    const SubtreeMerge f1;
    const PrefixSuffixMerge f2;
    const Mapping f3;
    const T IE;     // Identity Element
    vector<Weight> weight;
    vector<T> subtree_dp, rerooting_dp;
    
    VertexWeightedReRooting(int V, const SubtreeMerge s, const PrefixSuffixMerge p, const Mapping m, const T I)
        : g(V), f1(s), f2(p), f3(m), IE(I), weight(V), subtree_dp(V, I), rerooting_dp(V, I) {}

    void add_edge(int u, int v){
        g[u].emplace_back(v);
        g[v].emplace_back(u);
    }

    void set(int v, Weight w){ weight[v] = w; }

    void subtree_dfs(int from, int par){
        subtree_dp[from] = f3(subtree_dp[from], weight[from]);
        for(int to : g[from]){
            if(to != par){
                subtree_dfs(to, from);
                subtree_dp[from] = f1(subtree_dp[from], subtree_dp[to]);
            }
        }
    }

    void rerooting_function(int cur, T par_val, int par){
        int V = g[cur].size();
        vector<T> prefix(V + 1, IE), suffix(V + 1, IE);
        for(int i = 0; i < V; ++i){
            int child = g[cur][i];
            if(child == par)    prefix[i + 1] = f1(prefix[i], par_val);
            else    prefix[i + 1] = f1(prefix[i], subtree_dp[child]);
        }
        for(int i = V - 1; i >= 0; --i){
            int child = g[cur][i];
            if(child == par)    suffix[i] = f1(suffix[i + 1], par_val);
            else    suffix[i] = f1(suffix[i + 1], subtree_dp[child]);
        }
        for(int i = 0; i < V; ++i){
            int child = g[cur][i];
            if(child == par){
                rerooting_dp[cur] = f1(subtree_dp[cur], par_val);
            }
            else{
                rerooting_function(child, f3(f2(prefix[i], suffix[i + 1]), weight[cur]), cur);
            }
        }
    }

    vector<T> build(int root = 0){
        subtree_dfs(root, -1);
        rerooting_function(root, IE, -1);
        rerooting_dp[root] = subtree_dp[root];
        return rerooting_dp;
    }
};

使用例 (ネタバレ注意)

重みなし

atcoder.jp

▶解法

subtree_dp[v] := v を黒で塗ったときの v 以下の部分木の塗り方とする。
これは、それぞれの子ノード child が白で塗る 通りまたは黒で塗る subtree_dp[child] 通りであるから、( subtree_dp[child]) の総積で求まる。
累積積は上の値の総積になるので、PrefixSuffixMerge には単なる積を定義しておく。

▶ソースコードを展開する

int main(){
    int n, m;
    cin >> n >> m;
    auto f1 = [&](ll a, ll b){return a * (1LL + b) % m;};
    auto f2 = [&](ll a, ll b){return a * b % m;};
    ReRooting<ll> g(n, f1, f2, 1LL);
    for(int i = 1; i < n; ++i){
        int x, y;
        cin >> x >> y;
        --x, --y;
        g.add_edge(x, y);
    }
    auto dp = g.build();
    for(int i = 0; i < n; ++i)  cout << dp[i] << '\n';
}

重み付き

http://judge.u-aizu.ac.jp/onlinejudge/description.jsp?id=1595
木の直径を求める。
2 回 DFS すればできるが、全方位木 DP でやってみる。

▶解法

subtree_dp[v] := 葉から v までの距離の最大値

▶ソースコードを展開する

int main(){
    int n;
    cin >> n;
    auto f1 = [](int a, int b){return max(a, b);};
    auto f2 = [](int a, int b){return max(a, b);};
    auto f3 = [](int a, int b){return a + b;};
    WeightedReRooting<int, int> g(n, f1, f2, f3, 0);
    for(int i = 1; i < n; ++i){
        int s, t, w;
        cin >> s >> t >> w;
        g.add_edge(s, t, w);
    }
    auto dp = g.build();
    int ans = 0;
    for(int i = 0; i < n; ++i)  ans = max(ans, dp[i]);
    cout << ans << '\n';
}

頂点重み付き

codeforces.com

▶解法

Staircase Nim を拡張したような問題。

頂点からの距離 の累積 xor を持った 次元ベクトルを持たせる。
勝敗判定は すべての xor を取ればよい。
これは、 の部分が勝敗に影響を及ぼさないからである。
例えば、現在の状態 (Grundy 数) が負けの状態であるときに をいじると勝ちに誘導できると仮定しても、そのままその手を返すという正しい手を相手に選択されると再び負けの状態になる。
よって、そもそも仮定が成り立たない。

こういう不偏ゲームの考え方苦手だ......

▶ソースコードを展開する

int main(){
    int n, k, u, v, a;
    cin >> n >> k;
    using Vec = vector<int>;
    auto f1 = [&](Vec l, Vec r){
        for(int i = 1; i < k * 2; ++i)  l[i] ^= r[i - 1];
        l[0] ^= r[k * 2 - 1];
        return l;
    };
    auto f2 = [&](Vec l, Vec r){
        for(int i = 0; i < k * 2; ++i)  l[i] ^= r[i];
        return l;
    };
    auto f3 = [&](Vec v, int x){
        v[0] ^= x;
        return v;
    };
    Vec I(k * 2, 0);
    VertexWeightedReRooting<Vec, int> g(n, f1, f2, f3, I);
    for(int i = 1; i < n; ++i){
        cin >> u >> v;
        --u, --v;
        g.add_edge(u, v);
    }
    for(int i = 0; i < n; ++i){
        cin >> a;
        g.set(i, a);
    }
    auto dp = g.build();
    for(int i = 0; i < n; ++i){
        int win = 0;
        for(int j = k; j < k * 2; ++j)  win ^= dp[i][j];
        cout << (win ? 1 : 0) << (i == n - 1 ? '\n' : ' ');
    }
}

参考文献

もうひとつの全方位木DP - ei1333の日記

雑記

全方位木 DP しばらく書かずにいるとバグらせがちなので、ライブラリにしてみた。
ただ、ライブラリを書き上げてから、もっと他にやるべきことがあるだろ？暇人かよ、という気持ちになり悲しくなった。

こういうのは GitHub に上げるのがいいのだろうか。

コードを見れば分かるのだが、効率の悪いことをしている。
自分は高速化にそこまで興味がないので放置してしまっているが、直してくれる誰かに期待......