BZOJ 3522 Hotel 完整解析

樹上求無序三元組 (x,y,z) 互不相同且 dis(x,y)=dis(y,z)=dis(x,z)。N ≤ 5000。

幾何洞察：樹上三點等距只有兩種結構

1. 共線：三點在一條鏈上 → 無解（距離無法全等）
2. 星形：存在唯一中心點，三點分別在三條不同分支上，且到中心距離相等

故只需枚舉中心點，統計各分支各深度節點數，組合計算。

演算法：枚舉中心點 + DFS 深度計數

對每個節點 u 作中心：

1. 對每條出邊 v，DFS 統計 cnt[d] = 深度 d 的節點數
2. 合併各分支：對每個深度 d，三元組數 += 選三條不同分支各取一個深度 d 節點

用前綴和優化合併：f[d] = 已處理分支深度 d 總數，g[d] = 兩分支組合數。新分支 cnt[d] 貢獻 g[d] * cnt[d]，更新 g[d] += f[d] * cnt[d]，f[d] += cnt[d]。

完整 AC 程式碼

#include <bits/stdc++.h>
using namespace std;

const int MAXN = 5050;
vector G[MAXN];
int n;
long long ans = 0;
int temp[MAXN], f[MAXN], g[MAXN];

void dfs(int u, int p, int dpt) {
    temp[dpt]++;
    for (int v : G[u])
        if (v != p) dfs(v, u, dpt + 1);
}

int main() {
    ios::sync_with_stdio(false);
    cin.tie(nullptr);
    
    cin >> n;
    for (int i = 1; i < n; i++) {
        int u, v; cin >> u >> v;
        G[u].push_back(v);
        G[v].push_back(u);
    }
    
    for (int u = 1; u <= n; u++) {
        fill(f, f + n + 1, 0);
        fill(g, g + n + 1, 0);
        
        for (int v : G[u]) {
            fill(temp, temp + n + 1, 0);
            dfs(v, u, 1);
            
            for (int d = 1; d <= n; d++) {
                ans += 1LL * g[d] * temp[d];
                g[d] += f[d] * temp[d];
                f[d] += temp[d];
            }
        }
    }
    
    cout << ans << '\n';
    return 0;
}

關鍵點總結

• 星形結構必要性：樹上三點等距 ⇔ 存在中心點，三點在不同分支同深度
• O(N²) 複雜度：N=5000 時 2.5×10⁷ 操作，C++ 極快
• 前綴和合併：f,g 陣列將 O(度²) 降為 O(度)
• 長整型：答案可能達 C(5000,3) ~ 2×10¹⁰，需 long long