題目鏈接

問題描述

給定一棵樹，樹中每個結點權值為[-100,100]之間的整數。樹中包含結點總數不超過1e5。任選兩個非根節點A、B，將這兩個結點與其父節點斷開，可以得到三棵子樹。現要求三棵子樹的權值之和相等，問A、B有多少種選擇方法。

輸入

T：樣例種數
N：樹中結點個數
v1 father1
v2 father2

問題分析

此問題是一道樹形DP。
如何才能將一棵樹劃分成三棵權值之和相等的子樹？有兩種情況：

• 若結點x和結點y沒有血緣關系（不是祖孫關系），則x和y的權值之和都是s（s為整棵樹的權值之和的三分之一），x和y把這棵樹截為三段。
• 若結點x和結點y有血緣關系，不妨設x是y的祖先，則x的權值之和為2s，y的權值之和為s。x和y把這棵樹截為三段。

最精簡的代碼

#include<iostream>
#include<stdio.h> 
#include<iostream> 
using namespace std;
const int maxn = 1e5 + 7;
typedef  long long ll;
struct Node {
    int v;
    int s; 
    int son;
}a[maxn];
int nex[maxn];
int root;
int per;
int n;
ll ans;
int perCount = 0;
int go(int nodeId) {
    int temp = perCount;
    a[nodeId].s = a[nodeId].v;
    for (int i = a[nodeId].son; i != -1; i = nex[i]) {
        a[nodeId].s += go(i);
    }
    if (nodeId != root) {
        if (a[nodeId].s == per * 2) {
            ans += perCount - temp;
        }
        if (a[nodeId].s == per) {
            ans += temp;
            perCount++;
        }
    } 
    return a[nodeId].s;
}
void push(int parent, int son) {
    int temp = a[parent].son;
    nex[son] = temp;
    a[parent].son = son;
} 
int main() { 
    int T;
    cin >> T;
    while (T-- > 0) { 
        cin >> n;
        for (int i = 0; i <= n; i++) {
            a[i].son = -1;
            nex[i] = -1;
        }
        int s = 0;
        for (int i = 0; i < n; i++) {
            int father;
            cin >> a[i].v >> father;
            s += a[i].v;
            father--;
            if (father == -1) {
                root = i;
            }
            else {
                push(father, i);
            }
        }
        
        if (s % 3 == 0) {
            per = s / 3;
            ans = 0;
            perCount = 0;
            go(root); 
            cout << ans << endl;
        }
        else {
            cout << 0 << endl;
        } 
    }
    return 0;
}
 

復雜但是直觀的代碼

#include<iostream>
#include<stdio.h>
#include<stdlib.h>
using namespace std;
const int maxn = 1e5 + 7;
typedef long long ll;
int n;
int per;//每個分支應該等於的數值
struct Node {
    int v;//結點權重
    int s;//結點所代表的子樹的權重之和
    int sonPerCount;//子樹中權值之和為per的結點個數（包括自身）
    int son;//兒子結點，鏈表第一個結點
    int next;//下一個兄弟結點
}a[maxn];
void push(int father, int son) {
    int temp = a[father].son;
    a[son].next = temp;
    a[father].son = son;
}
int root;
void init(int nodeId) {
    a[nodeId].s = a[nodeId].v;
    for (int i = a[nodeId].son; ~i; i = a[i].next) {
        init(i);
        a[nodeId].s += a[i].s;
        a[nodeId].sonPerCount += a[i].sonPerCount;
    }
    if (a[nodeId].s == per) {
        a[nodeId].sonPerCount++;
    }
}
ll count1 = 0, count2 = 0;//count1表示我的值為per時，count2表示我的值為2per時
int totalPer = 0;
int cnt = 0;
ll ans = 0;
void go(int nodeId) {
    if (a[nodeId].s == per) {
        if (nodeId != root) {
            count1 += totalPer - cnt - a[nodeId].sonPerCount;
        }
        cnt++;
    }
    //此處不能有else，因為當per=0時，子樹中的總和為per的結點也會生效
    if (a[nodeId].s == per * 2) {
        if (nodeId != root) {
            count2 += a[nodeId].sonPerCount;
            if (per == 0)count2--;//因為sonPerCount包括結點自身，所以需要先去掉結點
        }
    }
    for (int i = a[nodeId].son; ~i; i = a[i].next) {
        go(i);
    }
    if (a[nodeId].s == per)cnt--;
}
int main() {
    freopen("in.txt", "r", stdin);
    int T; cin >> T;
    while (T--) {
        cin >> n;
        int s = 0;
        for (int i = 1; i <= n; i++) {
            a[i].sonPerCount = 0;
            a[i].son = -1;
            a[i].next = -1;
        }
        for (int i = 1; i <= n; i++) {
            int father;
            cin >> a[i].v >> father;
            if (father == 0) {
                root = i;
            }
            else {
                push(father, i);
            }
            s += a[i].v;
        }
        if (s % 3 == 0) {
            per = s / 3;
            init(root);
            count1 = count2 = 0;
            totalPer = 0;
            for (int i = 1; i <= n; i++) {
                if (a[i].s == per)totalPer++;
            }
            go(root);
            ans = count1 / 2 + count2;
            cout << ans << endl;
        }
        else { cout << 0 << endl; }
    }
    return 0;
}
 

註意事項

• 如果用Java寫，此題會超時。因為輸入量比較大
• 如果數據充分一些，1e5的復雜度有可能爆棧，所以需要使用棧的方式來遍歷樹（也可以先對樹進行線索化）。

hihocoder第237周：三等分帶權樹