Algorithm

HackerRank/Algorithm/Dynamic Programming/Kingdom Division

Problem Summary

Given a tree with N nodes, find the number of ways to divide the nodes into two sets, such that every node has at least one node that is connected to it and in the same set with it.

Solution

This problem can be solved with DP on trees.

Let f[i,0] be the number of ways to divide the subtree with node i as root, into two sets, when i is in different set with its parent.
Similarly, let f[i,1] be the number of ways to divide the subtree when i is in the same set with its parent.

If node i has no child, then f[i,0]=0, f[i,1]=1.
If node i have M children, c[1],c[2],…,c[M],
f[i,1] = sum { f[c[1],k[1]] × f[c[2],k[2]] × … × f[c[M],k[M]]}, k[j] = 0,1.
f[i,0] = f[i,1] - f[c[1],0] × f[c[2],0] × … × f[c[M],0].

So we only need to choose a node whose degree is 1 as the root, and use DFS to calculate f. The answer is 2 × f[root,0], due to symmetry.

Code

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
#include <cstdio>
#include <iostream>
#include <cstring>
#include <cstdlib>
#include <vector>
using namespace std;
const int maxn = 100002;
const int maxm = 200002;
const long long mo = 1e9 + 7;
int n,degree[maxn];
bool vis[maxn];
int tot=0,head[maxm],next_e[maxm],v[maxm];
long long f[maxn][2];
void add_e(int a,int b)
{
    tot++; next_e[tot]=head[a]; head[a]=tot; v[tot]=b;
}
void dfs(int cur,long long sum,int x,vector<int> &ch)
{
    if(cur==ch.size())
    {
        f[x][1] = (f[x][1]+sum)%mo;
        return;
    }
    for(int j=0;j<=1;j++)
    {
        if(cur==0) dfs(cur+1,f[ch[cur]][j],x,ch);
        else dfs(cur+1,(sum*f[ch[cur]][j])%mo,x,ch);
    }
}
void calc(int x)
{
    vis[x]=true;
    std::vector<int> children;
    for(int i = head[x]; i!=0 ;i = next_e[i])
    {
        int y = v[i];
        if(vis[y]) continue;
        children.push_back(y);
        calc(y);
    }
    if(children.size()==0)
    {
        f[x][0]=0;
        f[x][1]=1;
        return;
    }
    if(children.size()==1)
    {
        f[x][0]=f[children[0]][1];
        f[x][1]=(f[x][0]+f[children[0]][0])%mo;
        return;
    }
    dfs(0,0,x,children);
    long long tmp = f[children[0]][0];
    for(int i=1;i<children.size();i++) 
        tmp = tmp * f[children[i]][0] % mo;
    
    f[x][0] = f[x][1] + mo - tmp;
    if(f[x][0]>=mo) f[x][0] -= mo;
}
int main()
{
    scanf("%d",&n);
    int a,b;
    for(int i=1;i<n;i++)
    {
        scanf("%d%d",&a,&b);
        add_e(a,b);
        add_e(b,a);
        degree[a]++;
        degree[b]++;
    }
    int root = 0;
    for(int i=1;i<=n;i++)
        if(degree[i]==1)
        {
            root=i;
            break;
        }
    memset(vis,0,sizeof(vis));
    calc(root);
    long long ans = f[root][0]<<1;
    if(ans>=mo) ans-=mo;
    printf("%lld\n",ans);
    return 0;
}