題目連結

題目描述

給定一顆n個點n−1條邊的樹，每條邊的長度都是1，i號節點的權是\(a_i\)。如果三個節點兩兩在樹上的距離都是4，那麼稱這三個節點構成了一個“組合”，一個“組合”的權是三個節點的權的乘積。求所有“組合”的權的和。

資料範圍：$$1\le n\le 10^{5} , 1\le a_i \le 10$$

題目分析

• 性質1

  直覺上，一個組合中的點一定是這種形狀（紅色點）：
  

下面證明：
設三個點為A,B,C。記d(x,y)為x到y的距離,記A到B與B到C的路徑的重疊部分長為t。
由於樹上兩點間簡單路徑是唯一的，所以$$d(A,C)=d(A,B)+d(B,C)-2t$$
$$t=2$$
Q.E.D

• 性質2

  定義上圖中黃點為組合的中心點。
  顯然 **以一個點為中心點的所有組合的權值之和** 即 **到此點距離為2的點的集合 的 三元子集 的權值之和**
  

• 性質3

  根據[這道題](https://www.luogu.com.cn/problem/U138272),我們有O(n)的方法求出一個集合的三元子集的權值之和
  

程式碼（沒過）

#include<iostream>
#include<queue>
#define MAXN 100005
using namespace std;
int n;
long long ans;
int pw[MAXN],head[MAXN],to[MAXN<<1],nxt[MAXN<<1],cnt;
inline void add(int a,int b){
    ++cnt;
    to[cnt]=b;
    nxt[cnt]=head[a];
    head[a]=cnt;
    return ;
}
inline int read(){
    int x=0;char c;
    c=getchar();
    while(c<'0'||c>'9') c=getchar();
    while(c>='0'&&c<='9'){
        x=x*10+c-'0';
        c=getchar();
    }
    return x;
}
queue<int> q1,q2;
int s[4][MAXN];
int main(){
    n=read();
    for(int i=1;i<=n;++i) pw[i]=read();
    int a,b;
    for(int i=1;i<n;++i){
        a=read(),b=read();
        add(a,b);
        add(b,a);
    }
    int c,now;
    for(int i=1;i<=n;++i){
        while(!q2.empty()) q2.pop();
        for(int j=head[i];j;j=nxt[j]) q1.push(to[j]);
        while(!q1.empty()){
            now=q1.front();q1.pop();
            for(int j=head[now];j;j=nxt[j]){
                if(to[j]==i) continue;
                q2.push(to[j]);
                //cout<<"!"<<to[j]<<endl;
            }
            //cout<<"!"<<now<<endl;
        }
        for(c=1;!q2.empty();++c){
            s[0][c]=pw[q2.front()];q2.pop();
            //cout<<"?"<<s[0][i]<<endl;
        }
        --c;
        s[1][0]=s[2][0]=s[3][0]=0;
        for(int i=1;i<=c;++i) s[1][i]=s[1][i-1]+s[0][i]; //一元組的和 字首和
        for(int i=1;i<=c;++i) s[2][i]=s[2][i-1]+s[1][i-1]*s[0][i]; //二元組的和
        for(int i=1;i<=c;++i) s[3][i]=s[3][i-1]+s[2][i-1]*s[0][i];
        //ans+=/*q2的三元組的和*/;
        ans+=s[3][c];
        //cout<<ans<<' '<<c<<endl;
    }
    cout<<ans;
    return 0;
}