【洛谷5002】ひたむきOI-先祖探し(樹上計数)

OI-祖先を探す

テーマの背景

(Imakf)は小さなコンニャクで、最近習ったばかりで、携帯電話(APP)でゲームを見て(LCA)とダウンロードしました.

タイトルの説明

このゲームでは、(N)個のノードがあり、ルートノードが(R)、システムが(M)個の点(P_1,P_2...P_M)を選択し、(Imakf)に何組の点対((u_i,v_i))があるかを答える最近の共通の祖先が(P_i)です.(Imakf)こんにゃくなので、習ってもできないので、助けを求めるしかありません.
(Imakf)結局少し習ったことがあるので、彼はあなたに答えの型((10^9+7))を許可します.

にゅうしゅつりょくけいしき

入力形式:

最初の行(N,R,M).
その後(N-1)行ごとに2つの数(a,b)は(a,b)の間に1つのエッジを表します.
その後(1)行(M)個数表示(P_i)

出力フォーマット:

(M)行、行ごとに1つの数、ii行目の数は、点対((u_i,v_i))の最近の共通祖先がどれだけあるかを示す(P_i)

入出力サンプル

サンプル#1を入力:

7 1 3
1 2
1 3
2 4
2 5
3 6
3 7
1 2 4

出力サンプル#1:

31
7
1

説明

問い合わせ1:
(1,1)(1,2)(1,3)(1,4)(1,5)(1,6)(1,7)(2,1)(2,3)(2,7)(3,1)(3,2)(3,4)(3,5)(4,3)(4,3)(4,6)(4,7)(5,1)(5,6)(5,6)(5,6)(5,7)(5,6)(6,7)(6,2)(6,4)(6,4,5)(7,2)(7,4)(7,4)(7,4)(7,4)(7,4)(7,4)(7,4)(7,4)(7,4)(7,4)(7,4)(7,5)(7,5)計31組
問い合わせ2:
(2,2)(2,4)(2,5)(4,2)(4,5)(5,2)(5,4)計7群
質問3:
(4,4)合計1組

データ範囲

\(N\leq10000,M\leq50000\)

問題解

これが洛谷の新しい問題だと見て、少し考えてやった.
私は(dfs)の時に(ans)配列を前処理しました.つまり、各点の答えなので、時間の複雑さは(max(n,m))です.
各ポイント(x)について、答えを(x)をまたぐシナリオと、(x)をまたぐシナリオの2つの部分に分けることができます.

xを越えていない部分:

明らかに、その中の1つの点は必ず(x)であるため、この部分の答えは(x)の(siz)サイズの(2)倍(-1)(点対であり(1,1)が1つであるため)である.
すなわち(ans 1=2*siz[x]-1)である.

(x)にまたがる部分:

乗算原理によれば、すべての(x)の息子(son_i)の(siz)サイズを乗算することがこの部分の答えである.
(x)全部で(k)本の木があるとすると[ans 2=sum_{i=1}^{k}sum_{j=1}^{k}siz[son[i]*siz[son[j]]\=sum_{i=1}^{k}siz[son[i]*(siz[x]-1)\=(siz[x]-1)*(siz[x]-1)]しかし、私たちは一部を重く計算したので、((i=j))の部分ではないことに気づきます.
だから私たちは(sum_{i=1}^{k}siz[son[i]^2)を減算します.
だから(ans 2=(siz[x]-1)^2-sum_{i=1}^{k}siz[son[i]]^2\)

総合:

\[ Ans=ans1+ans2\\=2*siz[x]-1+(siz[x]-1)^2-\sum_{i=1}^{k}siz[son[i]]^2\\=siz[x]^2-\sum_{i=1}^{k}siz[son[i]]^2\]
各サブツリーの(siz[son[i]^2)の和を維持するだけで答えを得ることができます.

ツッコミ一言:この问题Ans最大は(N^2)で、意外にも型を取りに行きます。。。

code：

#include
#include
#include
#include
#include
#define N 10005
#define R register
using namespace std;
templateinline void read(T &a){
    char c=getchar();T x=0,f=1;
    while(!isdigit(c)){if(c=='-')f=-1;c=getchar();}
    while(isdigit(c)){x=(x<<1)+(x<<3)+c-'0';c=getchar();}
    a=f*x;
}
int n,rt,m,tot,p;
int siz[N],ans[N],h[N],sum[N];
struct node{
    int nex,to;
}edge[N<<1];
inline void add(R int u,R int v){
    edge[++tot].nex=h[u];
    edge[tot].to=v;
    h[u]=tot;
}
inline void dfs(R int x,R int f){
    siz[x]=1;
    for(R int i=h[x];i;i=edge[i].nex){
        R int xx=edge[i].to;
        if(xx==f)continue;
        dfs(xx,x);
        siz[x]+=siz[xx];
        sum[x]+=siz[xx]*siz[xx];
    }
    ans[x]=siz[x]*siz[x]-sum[x];
}
int main(){
    read(n);read(rt);read(m);
    for(R int i=1,u,v;i<=n-1;i++)
        read(u),read(v),add(u,v),add(v,u);
    dfs(rt,0);
    while(m--){
        read(p);
        printf("%d
",ans[p]);
    }
    return 0;
}