裏紙

ほぼ競プロ、たまに日記

TCO 2015 Round 1B Hard - TheTreeAndMan

問題

TopCoder Statistics - Problem Statement

問題概要

頂点数がNの根付き木を考える。そして、全ての辺は根から離れていく方向に向いている(つまり、どのN個の頂点にも根からたどり着くことが可能)。頂点をそれぞれ0~N-1で番号をふる。0の頂点を根とする。また、配列型で各頂点i(1 \le i \le N-1)に対して、その頂点の親が何であるかがparent_{i-1}で与えられる。

そして、Examples内にある図の棒人間みたいな形を"Man"と呼ぶことにする。これは"Man"の中でも最小単位のもので、腕が片方伸びたりしているものも"Man"になる。(7つの頂点組を考え、特定の頂点から特定の頂点へ到達可能(このルートは到達可能なら一意に定まる)であることと、その頂点組を構成した時に辺を共有しないことになる)。

与えられるグラフ内に"Man"がいくつあるかを数え上げて、10^{9} + 7で割った余りを答えよ。

イデア

Editorialをみながらまとめる。

まずは、頂点iを根とする部分木を考えて、その頂点iから到達可能な頂点数をreach_iに保存しておく。これは、その親の頂点が自身の頂点より必ず番号が小さいので、番号が大きい方からDP的に埋めていくことが出来る。

そして、次にleg_iを計算していきたい。この値の意味するところは、頂点iを"Man"の足の付根とするときに、足の2箇所の選び方が何通りあるかという情報を保持する。

まず、これを計算していくために、\displaystyle cur=\sum_{j} reach_jという値を計算しておく。ここで、jは頂点iをまさに親として持つ頂点のみを含む(つまり、最初はcurにはreach_iからi自身を除いた数が入っている...(?))。そして、これらの各jに対して、順に以下のような操作をしていく:

  • leg_i(cur-reach_j )*reach_jを加える。これによって、パスが被らないように足の位置を選ぶことが出来る(まだ見ていない自分以外の足の選び方*いま注目している根jの下にある中から足を選ぶ方法)

  • curからreach_jを除く(curの更新)

そして、次にbody_iの計算をしていこう。ここには、iを頂点とする部分木に含まれる頂点jに対して\displaystyle body_i =\sum_{j} leg_jとするものである。これの意味するところは、次で説明する。

最後に、このiを固定し、body_iに対して考えていく。頂点iの親をjとし、jをまさに体の中心(頭と腕2本の付け根)とすると、body_iにはjを体の中心としたときの足の選び方の総数が入っていると解釈できる。頭の位置の選び方は、頂点jの親以上の存在しかなり得ないので、頂点jの深さを事前に計算しておけばよいことになる。そして、腕の選び方としては、まず腕の位置を2本決めるのだが、これはまさに脚の決め方と同じであるのでleg_jで表される。しかし、脚に到達するパスとは被ってはいけないので、そのような組み合わせを取り除く必要がある。それは、iの部分木から1点と、それ以外の部分木から1点を選ぶような組み合わせである。それは、reach_i * (reach_j -1 -reach_i )で表される。以上より、

\displaystyle ans = \sum_{i} (depth_i - 1) * body_i * (leg_j - reach_i * (reach_j -1 -reach_i ) )の剰余を計算すれば良い。

最終的に1つの頂点を固定することで計算することがが可能になるのが驚き。後は、腕の計算に脚の前処理をそのまま利用していく方法は感動した。 他にも、体の中心をx、足の付根をyと固定してxからyに到達可能であるときに同様に部分木の被りをうまく除きながら掛け算していく方法(O(N^2)で間に合う方法)の方が、直感的にはわかりやすかった。

実装(C++)

class TheTreeAndMan {
    public:
    const long mod=1e9+7;

    int sum(int x, int y)
    {
        long ret=x;
        ret+=y;
        while(ret<0) ret+=mod;
        ret%=mod;
        return (int)ret;
    }

    int mul(int x, int y)
    {
        long ret=x;
        ret*=y;
        ret%=mod;
        return (int)ret;
    }

    int solve(vector<int> parent) {
        //頂点数
        int n=parent.size()+1;
        //グラフ
        vector<int> G[2000];

        //グラフ内でのその頂点の深さ
        int depth[2000]={0};
        //その頂点から到達可能な頂点数(自身も含む)
        int reach[2000]={0};
        //その頂点を足の付根とするときの選び方の総数
        int leg[2000]={0};
        //その頂点以下のlegの和
        int body[2000]={0};

        //有向辺を張る
        rep(i,parent.size()) G[parent[i]].pb(i+1);

        //頂点iから到達可能な頂点数を計算
        for(int i=n-1; i>=0; --i)
        {
            //まず自分自身を計上
            reach[i]=1;
            rep(j,G[i].size()) reach[i]=sum(reach[i],reach[G[i][j]]);
        }

        //BFSして深さを計算
        queue<int> que;
        que.push(0);
        int vis[2000]={0};
        vis[0]=1;
        depth[0]=0;
        while(!que.empty())
        {
            int v=que.front();
            que.pop();
            rep(i,G[v].size())
            {
                int nx=G[v][i];
                if(!vis[nx])
                {
                    vis[nx]=1;
                    depth[nx]=depth[v]+1;
                    que.push(nx);
                }
            }
        }

        //legの計算
        rep(i,n)
        {
            int cur=0;
            rep(j,G[i].size()) cur=sum(cur,reach[G[i][j]]);

            rep(j,G[i].size())
            {
                leg[i]=sum(leg[i],mul(reach[G[i][j]],sum(cur,-reach[G[i][j]])));
                cur=sum(cur,-reach[G[i][j]]);
            }

            body[i]=leg[i];
        }

        //bodyの計算
        for(int i=n-1; i>0; --i)
        {
            body[parent[i-1]]=sum(body[parent[i-1]],body[i]);
        }

        int ret=0;
        for(int i=1; i<n; ++i)
        {
            int j=parent[i-1];
            int add=mul(depth[i]-1,mul(body[i],sum(leg[j],-mul(reach[i],reach[j]-1-reach[i]))));

            ret=sum(ret,add);
        }

        return ret;
    }
};