[Codeforces 161D] Distance in Tree

[题目链接]

         https://codeforces.com/contest/161/problem/D

[算法]

        点分治

        记cnt[u][i]表示以u为根的子树中深度为i的点有多少个

        时间复杂度 : O(NlogN)

[代码]

        

#include<bits/stdc++.h>
using namespace std;
#define MAXN 50010
#define MAXK 510
const int inf = 1e9;
typedef long long ll;
typedef long double ld;
typedef unsigned long long ull;

struct edge
{
        int to , nxt;
} e[MAXN << 1];

int n , k , tot , root;
ll ans;
int weight[MAXN] , head[MAXN] , size[MAXN];
bool visited[MAXN];
int cnt[MAXN][MAXK];

template <typename T> inline void chkmax(T &x,T y) { x = max(x,y); }
template <typename T> inline void chkmin(T &x,T y) { x = min(x,y); }
template <typename T> inline void read(T &x)
{
    T f = 1; x = 0;
    char c = getchar();
    for (; !isdigit(c); c = getchar()) if (c == '-') f = -f;
    for (; isdigit(c); c = getchar()) x = (x << 3) + (x << 1) + c - '0';
    x *= f;
}
inline void addedge(int u , int v)
{
        ++tot;
        e[tot] = (edge){v , head[u]};
        head[u] = tot;
}
inline void getroot(int u , int par , int total)
{
        size[u] = 1;
        weight[u] = 0;
        for (int i = head[u]; i; i = e[i].nxt)
        {
                int v = e[i].to;
                if (v == par || visited[v]) continue;
                getroot(v , u , total);
                size[u] += size[v];
                chkmax(weight[u] , size[v]);         
        }        
        chkmax(weight[u] , total - size[u]);
        if (weight[u] < weight[root]) root = u;
}
inline void calc(int u , int par , int depth)
{
        if (depth > k) return;
        ans += cnt[root][k - depth];
        for (int i = head[u]; i; i = e[i].nxt)
        {
                int v = e[i].to;
                if (visited[v] || v == par) continue;
                calc(v , u , depth + 1);
        }
}
inline void update(int u , int par , int depth)
{
        if (depth > k) return;
        ++cnt[root][depth];
        for (int i = head[u]; i; i = e[i].nxt)
        {
                int v = e[i].to;
                if (visited[v] || v == par) continue;
                update(v , u , depth + 1);
        }
}
inline void work(int u)
{
        visited[u] = true;
        cnt[u][0] = 1;
        for (int i = head[u]; i; i = e[i].nxt)        
        {
                int v = e[i].to;
                if (visited[v]) continue;
                calc(v , u , 1);
                update(v , u , 1);
        }
        for (int i = head[u]; i; i = e[i].nxt)
        {
                int v = e[i].to;
                if (visited[v]) continue;
                root = 0;
                getroot(v , u , size[v]);
                work(root);
        }
}

int main()
{
        
        read(n); read(k);
        for (int i = 1; i < n; i++)
        {
                int x , y;
                read(x); read(y);
                addedge(x , y);
                addedge(y , x);
        }
        weight[root = 0] = inf;
        getroot(1 , 0 , n);
        work(root);
        printf("%I64d\n" , ans);
        
        return 0;
    
}