Acwing-252-树(点分治)

链接:

https://www.acwing.com/problem/content/254/

题意:

给定一个有N个点（编号0,1,…,N-1）的树，每条边都有一个权值（不超过1000）。

树上两个节点x与y之间的路径长度就是路径上各条边的权值之和。

求长度不超过K的路径有多少条。

思路:

点分治, 就是将一棵树根据他的重心分成多颗子树求解.

代码:

#include 
using namespace std;
typedef long long LL;
const int MAXN = 1e4+10;
const int INF = 1e9;

struct Edge
{
    int to;
    int dis;
};
vector G[MAXN];
int Size[MAXN], Cnt[MAXN];
int Vis[MAXN], Dis[MAXN];
int root, maxroot;
int n, k, cnt, treesize;
LL ans = 0;

void GetRoot(int u, int fa)
{
    Size[u] = 1;
    int num = 0;
    for (int i = 0;i < G[u].size();i++)
    {
        int to = G[u][i].to;
        if (to == fa || Vis[to] == 1)
            continue;
        GetRoot(to, u);
        Size[u] += Size[to];
        num = max(num, Size[to]);
    }
    num = max(num, treesize-Size[u]);
    if (num < maxroot)
    {
        root = u;
        maxroot = num;
    }
}

void GetDis(int u, int fa)
{
    Cnt[++cnt] = Dis[u];
    for (int i = 0;i < G[u].size();i++)
    {
        int to = G[u][i].to;
        if (to == fa || Vis[to] == 1)
            continue;
        Dis[to] = Dis[u] + G[u][i].dis;
        GetDis(to, u);
    }
}

LL Cal(int u, int val)
{
    cnt = 0;
    Dis[u] = val;
    GetDis(u, 0);
    LL sum = 0;
    int l = 1, r = cnt;
    sort(Cnt+1, Cnt+1+cnt);
    while (l < r)
    {
        if (Cnt[l]+Cnt[r] <= k)
            sum += r-l, ++l;
        else
            --r;
    }
    return sum;
}

void Dfs(int u)
{
    ans += Cal(u, 0);
    Vis[u] = 1;
    for (int i = 0;i < G[u].size();i++)
    {
        int to = G[u][i].to;
        if (Vis[to] == 1)
            continue;
        ans -= Cal(to, G[u][i].dis);
        treesize = Size[to];
        maxroot = INF;
        GetRoot(to, 0);
        Dfs(root);
    }
}

int main()
{
    while(~scanf("%d%d", &n, &k) && (n||k))
    {
        ans = 0;
        for (int i = 1;i <= n;i++)
            G[i].clear();
        memset(Vis, 0, sizeof(Vis));
        memset(Size, 0, sizeof(Size));
        memset(Cnt, 0, sizeof(Cnt));
        int u, v, w;
        for (int i = 1;i < n;i++)
        {
            scanf("%d%d%d", &u, &v, &w);
            u++, v++;
//            cout << u << ' ' << v << ' ' << w << endl;
            G[u].push_back(Edge{v, w});
            G[v].push_back(Edge{u, w});
        }
        treesize = n;
        maxroot = INF;
        GetRoot(1, 0);
        Dfs(root);
        printf("%lld\n", ans);
    }

    return 0;
}