:warning: 全方位木DP
(graph_tree/reroot.hpp)

Code

#pragma once
#include<vector>

/**
 * @brief 全方位木DP
 */

template<typename T,typename F,typename Fix>
struct reroot{
    std::vector<std::vector<long long>>g;
    std::vector<int>p_list;
    std::vector<T>p_table;
    std::vector<bool>p_checked;
    std::vector<map<int,T>>table;
    std::vector<T>ans;
    F f;
    Fix fix;
    reroot(const std::vector<std::vector<long long>>& g,F f=F(),Fix fix=Fix()):g(g),f(f),fix(fix){
        int n=g.size();
        p_list.resize(n,-1);
        p_checked.resize(n,0);
        table.resize(n);
        p_table.resize(n,e);
        ans.resize(n,e);
        dfs1(0,-1);
        for(int i=0;i<n;++i)ans[i]=dfs2(i,-1);
    }
    T dfs1(int n,int p){
        p_list[n]=p;
        T tmp1=e,tmp2=e;
        std::vector<T>tmp(g[n].size());
        rep(i,g[n].size()){
            int t=g[n][i];
            if(t==p)continue;
            table[n][t]=tmp1;
            tmp1=f(tmp1,tmp[i]=dfs1(t,n));
        }
        for(int i=g[n].size()-1;i>=0;--i){
            int t=g[n][i];
            if(t==p)continue;
            table[n][t]=f(table[n][t],tmp2);
            tmp2=f(tmp[i],tmp2);
        }
        return fix(table[n][p]=tmp1,n,p);
    }
    T dfs2(int n,int p){
        if(n==-1)return e;
        if(!p_checked[n]){
            p_checked[n]=1;
            p_table[n]=dfs2(p_list[n],n);
        }
        if(p==-1){
            return f(table[n][p_list[n]],p_table[n]);
        }else{
            if(p_list[n]==-1)return fix(table[n][p],n,p);
            else return fix(f(table[n][p],p_table[n]),n,p);
        }
    }
    vector<T>query(){
        return ans;
    }
};
#line 2 "graph_tree/reroot.hpp"
#include<vector>

/**
 * @brief 全方位木DP
 */

template<typename T,typename F,typename Fix>
struct reroot{
    std::vector<std::vector<long long>>g;
    std::vector<int>p_list;
    std::vector<T>p_table;
    std::vector<bool>p_checked;
    std::vector<map<int,T>>table;
    std::vector<T>ans;
    F f;
    Fix fix;
    reroot(const std::vector<std::vector<long long>>& g,F f=F(),Fix fix=Fix()):g(g),f(f),fix(fix){
        int n=g.size();
        p_list.resize(n,-1);
        p_checked.resize(n,0);
        table.resize(n);
        p_table.resize(n,e);
        ans.resize(n,e);
        dfs1(0,-1);
        for(int i=0;i<n;++i)ans[i]=dfs2(i,-1);
    }
    T dfs1(int n,int p){
        p_list[n]=p;
        T tmp1=e,tmp2=e;
        std::vector<T>tmp(g[n].size());
        rep(i,g[n].size()){
            int t=g[n][i];
            if(t==p)continue;
            table[n][t]=tmp1;
            tmp1=f(tmp1,tmp[i]=dfs1(t,n));
        }
        for(int i=g[n].size()-1;i>=0;--i){
            int t=g[n][i];
            if(t==p)continue;
            table[n][t]=f(table[n][t],tmp2);
            tmp2=f(tmp[i],tmp2);
        }
        return fix(table[n][p]=tmp1,n,p);
    }
    T dfs2(int n,int p){
        if(n==-1)return e;
        if(!p_checked[n]){
            p_checked[n]=1;
            p_table[n]=dfs2(p_list[n],n);
        }
        if(p==-1){
            return f(table[n][p_list[n]],p_table[n]);
        }else{
            if(p_list[n]==-1)return fix(table[n][p],n,p);
            else return fix(f(table[n][p],p_table[n]),n,p);
        }
    }
    vector<T>query(){
        return ans;
    }
};
Back to top page