graph/rerooting.cpp

• View this file on GitHub
• Last update: 2022-08-27 12:40:11+09:00

Code

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

struct DP {
    int dp;
    DP(int dp = 1) : dp(dp) {}
    DP operator+(const DP &a) const { // merge function
        return DP(dp + a.dp);
    }
    DP addroot() const { // g = g(f(dp_0,dp_1,...,dp_N),v) : add root
        return DP(dp + 1);
    }
};
DP unit = DP(); // unit

template <class T>
struct Rerooting {
    int N;
    vector<int> parents;
    vector<vector<int>> G;
    vector<vector<T>> dp;
    vector<T> ans;
    Rerooting(int N, const vector<vector<int>> &G) : N(N), G(G) {
        parents.assign(N, -1);
        dp.resize(N);
        ans.resize(N);
        dfs(0);
        bfs(0);
    }
    T dfs(int v, int p = -1) {
        T ret = unit;
        dp[v] = vector<T>(G[v].size());
        for (int i = 0; i < (int)G[v].size(); ++i) {
            if (p == G[v][i]) {
                parents[v] = i;
                continue;
            }
            dp[v][i] = dfs(G[v][i], v);
            ret      = ret + dp[v][i];
        }
        return ret.addroot();
    }
    void bfs(int v, const T &res_p = unit, int p = -1) {
        if (p != -1) dp[v][parents[v]] = res_p;
        int deg = G[v].size();
        vector<T> dpl(deg + 1), dpr(deg + 1);
        dpl[0]   = unit;
        dpr[deg] = unit;
        for (int i = 0; i < deg; ++i)
            dpl[i + 1] = dpl[i] + dp[v][i];
        for (int i = deg - 1; i >= 0; --i)
            dpr[i] = dpr[i + 1] + dp[v][i];
        ans[v] = dpr[0].addroot();
        for (int i = 0; i < deg; ++i) {
            if (parents[v] == i) continue;
            T d = dpl[i] + dpr[i + 1];
            bfs(G[v][i], d.addroot(), v);
        }
    }
    T operator[](int k) {
        return ans[k];
    }
};

#line 2 "graph/rerooting.cpp"
#include <bits/stdc++.h>
using namespace std;

struct DP {
    int dp;
    DP(int dp = 1) : dp(dp) {}
    DP operator+(const DP &a) const { // merge function
        return DP(dp + a.dp);
    }
    DP addroot() const { // g = g(f(dp_0,dp_1,...,dp_N),v) : add root
        return DP(dp + 1);
    }
};
DP unit = DP(); // unit

template <class T>
struct Rerooting {
    int N;
    vector<int> parents;
    vector<vector<int>> G;
    vector<vector<T>> dp;
    vector<T> ans;
    Rerooting(int N, const vector<vector<int>> &G) : N(N), G(G) {
        parents.assign(N, -1);
        dp.resize(N);
        ans.resize(N);
        dfs(0);
        bfs(0);
    }
    T dfs(int v, int p = -1) {
        T ret = unit;
        dp[v] = vector<T>(G[v].size());
        for (int i = 0; i < (int)G[v].size(); ++i) {
            if (p == G[v][i]) {
                parents[v] = i;
                continue;
            }
            dp[v][i] = dfs(G[v][i], v);
            ret      = ret + dp[v][i];
        }
        return ret.addroot();
    }
    void bfs(int v, const T &res_p = unit, int p = -1) {
        if (p != -1) dp[v][parents[v]] = res_p;
        int deg = G[v].size();
        vector<T> dpl(deg + 1), dpr(deg + 1);
        dpl[0]   = unit;
        dpr[deg] = unit;
        for (int i = 0; i < deg; ++i)
            dpl[i + 1] = dpl[i] + dp[v][i];
        for (int i = deg - 1; i >= 0; --i)
            dpr[i] = dpr[i + 1] + dp[v][i];
        ans[v] = dpr[0].addroot();
        for (int i = 0; i < deg; ++i) {
            if (parents[v] == i) continue;
            T d = dpl[i] + dpr[i + 1];
            bfs(G[v][i], d.addroot(), v);
        }
    }
    T operator[](int k) {
        return ans[k];
    }
};

Back to top page