裏紙

ほぼ競プロ、たまに日記

Codechef Nomvember Cook-Off 2017 - Adjacent leaves

問題

Contest Page | CodeChef

問題概要

根付き木の美しさを、good subsets of leavesの個数で定義する。

どういうものがgood subsets of leavesになるかというと、根からdfsをしていき、訪れた順番に、そのノードが葉であればその頂点番号をリストに追加する。こうすると、訪れる順番が何通りか考えられるので、その分のリストが得られる(dfs sequenceと呼ぶことにする)。そしてどのようなsubsetがgoodになるかというと、要素が1つ以上であり、そのsubsetがリストのどれかの連続部分列として現れてくれば良い。

頂点数nの木が与えられるので、各i (1 \le i \le n)をに対して、頂点iを根にした場合の根付き木の美しさを10^9+7で割った余りを答えよ。

  •  2 \le n \le 300000

例えば以下のような木がある:

f:id:imulan:20171222204241p:plain

1が根であるとすると、葉のノードは6個あるので、subsetの候補は2^6-1 = 63通りあるが、そのうちのsubsetでgoodでないものは以下の8通り:

{5,7,9}, {5,7,10}, {5,8,9} {5,8,10}, {6,7,9}, {6,7,10}, {6,8,9}, {6,8,10}

この3つの頂点をdfsで訪問順でみた時に連続で訪れることは有り得ない。good subsetの一例を挙げると、{5,7,9,10}などが挙げられる。これはdfsで出来る訪問リストが[6,5,9,10,7,8]となる時の真ん中の4つがこれに該当するのでgoodであると言える。

よって、上の木において根が1の場合は63-8 = 55通りになる。

イデア

single root

問題では各頂点について根を仮定した場合の答えを要求されるが、とりあえずは1つの根に対して答えを求めることを考えよう(ここでは根を1で固定する)。

以下のdpを考える:

  • dp[i][0] := iを根とする部分木に対するgood subsetの個数(ただし空のsubsetと、その部分木内にある全ての葉を含むようなsubsetは除く)
  • dp[i][1] := iを根とする部分木に対するgood subsetの中で、dfs sequenceのprefixとして現れるもの("good prefix subset"と呼ぶことにする)の個数(ただし空のsubsetと、その部分木内にある全ての葉を含むようなsubsetは除く)

今、頂点iのdpの値の更新を考える。

まず、good subsetの中で、そのsubsetの全てのLCAを取った時にそれがiの子孫になっているようなsubsetはそのLCAの位置で数え上げられているはずなので、それを持ち上げてくるだけで良い。なので、LCAiに一致するようなsubsetの個数を数えることを考える。

ここで、意識しなければならないのは、ある頂点iにいる時、dfs sequenceを作る時に、どの子に行くかを自由に選ぶことが出来るが、一度行く子を決めてしまうと、その子の部分木の頂点を全て訪問しきるまで次のiの子に行くということが出来ないということである。つまり、選んでくるsubsetをgoodなものにするためには、iの子の部分木内で中途半端に選んできて良いのはたかだか2種類までであり(それがprefixとsuffixになる)、それ以外の部分木に対しては全部の葉を取ってくるか1つも選ばないといけないことが分かる。

このことを踏まえて、dpの更新を考える。iの子のノードの個数をCとする。3パターンに分けて考える:

1つ目は、iの子の頂点集合の(空ではない)部分集合Sに対して、x \in Sの部分木に含まれる葉を全て取ってきたらそれはgood subsetになり得る。また、それはgood prefix subsetにもなり得るので、dp[i][0]とdp[i][1]の両方に加算される。具体的に加算される値は、(各部分木をまるまる選ぶか選ばないかなので)2^C通りあるが、dpの定義的に何も選ばないのと全部選ぶのは引いとく必要があるので2^C - 2となる。(このパートで、dpの定義で「その部分木内にある全ての葉を含むようなsubsetは除く」と書いたのが効いてくる。重複して数え上げてしまわないように出来ている。)

2つ目は、iの子の頂点uを1つ選び、u以外のiの子の部分集合Sを考える。x \in Sの部分木に含まれる葉を全て取り、u以下の部分木については全てを選ぶことはせずにgood prefix subsetとなるような葉の選び方(=dp[u][1])をすることを考える(すると、2つを組み合わせたsubsetもgood prefix subsetになる)。この場合のSの妥当な選び方は2^{C-1}-1通りになる(何も選ばないのは、LCAiの子孫になってしまい重複するのでダメ)。よって、これにdp[u][1]を掛け合わせたものをdp[i][0]とdp[i][1]の両方に加算すると良い。

3つ目は、iの子の頂点uv (ただし、u \neq v)とそれら以外のiの子の部分集合Sを考える。x \in Sの部分木に含まれる葉を全て取り、u以下の部分木とv以下の部分木については全てを選ぶことはせずにgood prefix subsetとなるような葉の選び方(dp[u][1],dp[v][1]通り)をすることを考えると、これらを先頭と末尾に配置することで、全体で(good prefix subsetにはならないが)good subsetを得ることができる。よって、これをdp[i][0]に加算していく。具体的に加算される値は、u \neq vに対して、dp[u][1]*dp[v][1]の和を計算したもの(全体から余計な部分を引くことを考えれば、O(n)で計算できる)に対し、2^{C-2}をかけたものになる。

すると、根を1にしたものに対しては、答えをdp[1][0] + 1で得ることが出来る(+1は、全ての葉を含んだものを数えている)。

all roots

全方位木dpの2回目のdfsで、全ての頂点に対してこの頂点を根にした時の答えを求めることを考えよう。今、根を1にした時の部分木の性質が得られている。dfsをしながら、今1ではない頂点にいる状況を考えて、根付き木上で子のノードの個数がC個で、親側の方にも子が1つあり、合計でC+1個の部分木があると見ることが出来る。子側のC個に対してはもう dp_0 , dp_1の値は得られていて、親側からもこの2値を伝播しながらdfsしていき、その結果を統合してそのノードに対する答えを得ることを考えたい。それは、親に対しても子と同じように dp_0 , dp_1の値を持ちながら、子に渡す時にこの2つの値を更新していくことで答えを求めることができる。この時に考えるべきことは、上で考えたことと同じことを考えれば良いが、遷移させる先の分まで数えてしまわないように注意。

実装(C++)

#include <bits/stdc++.h>
using namespace std;
using ll = long long;
#define rep(i,n) for(int (i)=0;(i)<(int)(n);++(i))
#define all(x) (x).begin(),(x).end()
#define pb push_back
#define fi first
#define se second

const ll mod = 1e9+7;
const int N = 300003;
ll pw[N];

inline ll mod_pow(ll x, ll n)
{
    ll ret = 1;
    while(n)
    {
        if(n&1) (ret*=x)%=mod;
        (x*=x)%=mod;
        n>>=1;
    }
    return ret;
}

inline ll mod_inv(ll x)
{
    return mod_pow(x,mod-2);
}

const ll MOD_INV = mod_inv(2);

vector<vector<int>> G;
vector<ll> ans;
vector<vector<ll>> dp;

void dfs(int x, int p)
{
    int C = 0;
    for(int nx:G[x])if(nx!=p)
    {
        dfs(nx,x);
        ++C;
        // children
        rep(i,2) (dp[x][i] += dp[nx][i]) %= mod;
    }

    // cond.1
    if(C>0)rep(i,2) (dp[x][i] += pw[C]-2+mod) %= mod;

    if(C>1)
    {
        ll s = 0;
        ll d = 0;
        for(int nx:G[x])if(nx!=p)
        {
            // cond.2
            rep(i,2) (dp[x][i] += dp[nx][1]*(pw[C-1]-1)) %= mod;

            (s += dp[nx][1]) %= mod;
            (d += dp[nx][1]*dp[nx][1]) %= mod;
        }

        // cond.3
        s = (s*s-d+mod)%mod;
        (s *= MOD_INV) %= mod;
        (s *= pw[C-2])%=mod;
        (dp[x][0] += s) %= mod;
    }
}

void dfs2(int x, int p, ll dp0, ll dp1)
{
    ans[x] = dp0+1;

    int C = (p!=-1);
    for(int nx:G[x])if(nx!=p)
    {
        ++C;
        // children
        (ans[x] += dp[nx][0]) %= mod;
    }

    // cond.1
    if(C>0) (ans[x] += pw[C]-2+mod) %= mod;

    if(C>1)
    {
        ll s = dp1;
        ll d = (dp1*dp1)%mod;

        // cond.2
        (ans[x] += dp1*(pw[C-1]-1)) %= mod;
        for(int nx:G[x])if(nx!=p)
        {
            // cond.2
            (ans[x] += dp[nx][1]*(pw[C-1]-1)) %= mod;

            (s += dp[nx][1]) %= mod;
            (d += dp[nx][1]*dp[nx][1]) %= mod;
        }

        // cond.3
        s = (s*s-d+mod)%mod;
        (s *= MOD_INV) %= mod;
        (s *= pw[C-2])%=mod;

        (ans[x] += s) %= mod;
    }

    // propagation to children
    ll nd0 = dp0, nd1 = dp1;
    for(int nx:G[x])if(nx!=p)
    {
        // children
        (nd0 += dp[nx][0]) %= mod;
        (nd1 += dp[nx][1]) %= mod;
    }
    C--; // 行き先のぶん

    // cond.1
    if(C>0)
    {
        (nd0 += pw[C]-2+mod) %= mod;
        (nd1 += pw[C]-2+mod) %= mod;
    }

    ll SUM = dp1;
    if(C>1)
    {
        ll s = dp1;
        ll d = (dp1*dp1)%mod;

        // cond.2
        (nd0 += dp1*(pw[C-1]-1)) %= mod;
        (nd1 += dp1*(pw[C-1]-1)) %= mod;
        for(int nx:G[x])if(nx!=p)
        {
            // cond.2
            (nd0 += dp[nx][1]*(pw[C-1]-1)) %= mod;
            (nd1 += dp[nx][1]*(pw[C-1]-1)) %= mod;

            (s += dp[nx][1]) %= mod;
            (d += dp[nx][1]*dp[nx][1]) %= mod;
            (SUM += dp[nx][1]) %= mod;
        }

        // cond.3
        s = (s*s-d+mod)%mod;
        (s *= MOD_INV) %= mod;
        (s *= pw[C-2])%=mod;

        (nd0 += s) %= mod;
    }

    // 子ノードにdfsを伸ばしていくが、その際に足してしまっている余計なものを引いた値を伝播
    for(int nx:G[x])if(nx!=p)
    {
        ll t0 = nd0, t1 = nd1;
        t0 = (t0 - dp[nx][0] + mod)%mod;
        t1 = (t1 - dp[nx][1] + mod)%mod;
        if(C>1)
        {
            t0 = (t0 - (dp[nx][1]*(pw[C-1]-1))%mod + mod)%mod;
            t1 = (t1 - (dp[nx][1]*(pw[C-1]-1))%mod + mod)%mod;

            ll dim = (dp[nx][1]*(SUM-dp[nx][1]))%mod;
            (dim *= pw[C-2])%=mod;
            t0 = (t0 - dim + mod)%mod;
        }
        dfs2(nx, x, t0, t1);
    }

    ans[x] %= mod;
}

void solve()
{
    int n;
    scanf(" %d", &n);

    G = vector<vector<int>>(n);
    ans = vector<ll>(n,0);
    dp = vector<vector<ll>>(n,vector<ll>(2,0));

    rep(i,n-1)
    {
        int u,v;
        scanf(" %d %d", &u, &v);
        --u; --v;
        G[u].pb(v);
        G[v].pb(u);
    }

    dfs(0,-1);
    dfs2(0,-1,0,0);
    rep(i,n) printf("%lld\n", ans[i]);
}

int main()
{
    pw[0] = 1;
    for(int i=1; i<N; ++i) pw[i] = (pw[i-1]*2)%mod;

    int T;
    scanf(" %d", &T);
    while(T--) solve();
    return 0;
}