点分治 nbut1654 Ancient battle tree

传送门：点击打开链接

题意：一棵树，求点对数，满足两点之间的最短距离等于树的最长链长。

思路：比较蠢，只想到点分治。

先用DFS求出树的最长链，然后跑点分治，这题很类似楼教主的男人八题的那题，不过那个更复杂些。

还有一个trick，就是如果n=1，只有1个点，那么最长链长等于0，满足的题意的点对有(0,0)，所以应该输出1，这里需要特判一下

#include<map>
#include<set>
#include<cmath>
#include<ctime>
#include<stack>
#include<queue>
#include<cstdio>
#include<cctype>
#include<string>
#include<vector>
#include<cstring>
#include<iostream>
#include<algorithm>
#include<functional>
#define fuck(x) cout<<"["<<x<<"]"
#define FIN freopen("input.txt","r",stdin)
#define FOUT freopen("output.txt","w+",stdout)
using namespace std;
typedef long long LL;

const int MX = 4e4 + 5;
const int MS = 4e4 + 5;
const int INF = 0x3f3f3f3f;

struct Edge {
int v, nxt, cost;
} E[MS];
int rear, Head[MX];
void edge_init() {
rear = 0;
memset(Head, -1, sizeof(Head));
}
void edge_add(int u, int v, int cost) {
E[rear].v = v;
E[rear].cost = cost;
E[rear].nxt = Head[u];
Head[u] = rear++;
}

int vis[MX], A[MX], cnt[MX];
int Tree_cnt(int u, int f) {
int ret = 1;
for(int i = Head[u]; ~i; i = E[i].nxt) {
int v = E[i].v;
if(v == f || vis[v]) continue;
ret += Tree_cnt(v, u);
}
return ret;
}

int G_DFS(int n, int u, int f, int &ansn, int &ansid) {
A[u] = 1; int ret = 0;
for(int i = Head[u]; ~i; i = E[i].nxt) {
int v = E[i].v;
if(v == f || vis[v]) continue;
int nxt = G_DFS(n, v, u, ansn, ansid);
ret = max(ret, nxt); A[u] += ret;
}
ret = max(ret, n - A[u]);
if(ret < ansn) ansn = ret, ansid = u;
return A[u];
}

int Tree_G(int u) {
int cnt = Tree_cnt(u, -1), ansn = INF, ansid = 0;
G_DFS(cnt, u, -1, ansn, ansid);
return ansid;
}

void Tree_deep(int u, int f, int d, int &DFN) {
A[++DFN] = d;
for(int i = Head[u]; ~i; i = E[i].nxt) {
int v = E[i].v, cost = E[i].cost;
if(v == f || vis[v]) continue;
Tree_deep(v, u, d + cost, DFN);
}
}

LL F_deal(int u, int k) {
if(k <= 0) return 0;
int DFN = 0; LL ans = 0;
for(int i = Head[u]; ~i; i = E[i].nxt) {
int v = E[i].v, cost = E[i].cost;
if(vis[v]) continue;
Tree_deep(v, u, cost, DFN);
}
for(int i = 1; i <= DFN; i++) cnt[i] = 0;
for(int i = 1; i <= DFN; i++) cnt[A[i]]++;
if(k <= DFN) ans += cnt[k];

for(int i = 1; i < (k + 1) / 2; i++) {
if(i <= DFN && k - i <= DFN) ans += (LL)cnt[i] * cnt[k - i];
}
if(k % 2 == 0 && k / 2 <= DFN) ans += (LL)cnt[k / 2] * (cnt[k / 2] - 1) / 2;
return ans;
}

LL solve(int u, int k) {
int g = Tree_G(u);
vis[g] = 1;

LL ans = F_deal(g, k);
for(int i = Head[g]; ~i; i = E[i].nxt) {
int v = E[i].v, cost = E[i].cost;
if(vis[v]) continue;

ans -= F_deal(v, k - 2 * cost);
ans += solve(v, k);
}
return ans;
}

int solve_line(int u, int from, int &ans) {
int Max1 = 0, Max2 = 0;
for(int id = Head[u]; ~id; id = E[id].nxt) {
int v = E[id].v;
if(v == from) continue;

int t = solve_line(v, u, ans) + 1;

if(t > Max1) {
Max2 = Max1;
Max1 = t;
} else if(t > Max2) Max2 = t;
}

ans = max(ans, Max1 + Max2);
return Max1;
}

int main() {
int n;
while(~scanf("%d", &n)) {
edge_init();
memset(vis, 0, sizeof(vis));
if(n == 1) {
printf("1\n");
continue;
}

for(int i = 1; i <= n - 1; i++) {
int u, v, cost;
scanf("%d%d", &u, &v);
edge_add(u, v, 1);
edge_add(v, u, 1);
}
int line = 0; solve_line(1, -1, line);
printf("%I64d\n", solve(1, line));
}
return 0;
}