题解 | 小红树上染色

栈优化的树形dp

#include<bits/stdc++.h>
using namespace std;
using ll=long long;
const ll N=1e5+5,mod=1e9+7;

vector<vector<ll>>adj(N);
vector<bool>vis(N,0);
vector<vector<ll>>dp(N,vector<ll>(2,1));


int main(){
    ios::sync_with_stdio(0),cin.tie(0),cout.tie(0);
    ll n;
    cin>>n;
    for(ll i=1;i<=n-1;i++){
        ll u,v;
        cin>>u>>v;
        adj[u].push_back(v);
        adj[v].push_back(u);
    }
    stack<pair<ll,bool>>st;
    st.push({1,0});
    vis[1]=1;

    while(!st.empty()){
        auto [u,ok]=st.top();
        st.pop();

        if(!ok){
            st.push({u,1});
            for(auto t:adj[u]){
                if(vis[t]==0){
                    st.push({t,0});
                    vis[t]=1;
                }
            }
        }else{
            vis[u]=0;
            for(auto t:adj[u]){
                if(vis[t]==0){
                    dp[u][0] = (dp[u][0] * dp[t][1]) % mod;
                    dp[u][1] = (dp[u][1] * (dp[t][0] + dp[t][1])) % mod; 
                }
            }
        }
    }
    cout << (dp[1][0] + dp[1][1]) % mod;

    return 0;
}