poj1741——树上的分治，树dp

最新推荐文章于 2021-10-20 23:16:56 发布

原创最新推荐文章于 2021-10-20 23:16:56 发布 · 665 阅读

0 ·

CC 4.0 BY-SA版权

本文介绍了一种求解树上所有路径长度不超过k的点对数量的算法。通过选取树的重心作为根节点，递归地统计每棵子树内符合条件的点对数，并考虑不同子树间的组合。

题意：求树上所有满足路径长度小于等于k的两点的对数。

思路还是很好理解的，对每一棵子树，以重心为根节点统计在不同子树且相互距离小于等于k的对数。就是实现的细节处理很多，主要都是递归处理的比较多。

#include <iostream>
#include <cstdio>
#include <cstring>
#include <algorithm>
#include <map>
#include <vector>
using namespace std;

const int maxn = 50000 + 10;
struct edge
{
    int to, dist, next;
}edges[maxn];
int edgehead[maxn], tot;
int n, k, vis[maxn], ans, root, num;
int mx[maxn], size[maxn], dis[maxn];

void addedge(int from, int to, int dist)
{
    edges[tot].dist = dist;
    edges[tot].to = to;
    edges[tot].next = edgehead[from];
    edgehead[from] = tot++;
}

void dfssize(int u, int fa)
{
    size[u] = 1;
    mx[u] = 0;
    for(int i = edgehead[u]; i != -1; i = edges[i].next)
    {
        int v = edges[i].to;
        if(v != fa && !vis[v])
        {
            dfssize(v, u);
            size[u] += size[v];
            mx[u] = max(mx[u], size[v]);
        }
    }
}

void dfsroot(int r, int u, int fa)
{
    mx[u] = max(mx[u], size[r] - size[u]);
    if(mx[u] < mx[root]) root = u;
    for(int i = edgehead[u]; i != -1; i = edges[i].next)
    {
        int v = edges[i].to;
        if(v != fa && !vis[v]) dfsroot(r, v, u);
    }
}

void dfsdis(int u, int d, int fa)
{
    dis[num++] = d;
    for(int i = edgehead[u]; i != -1; i = edges[i].next)
    {
        int v = edges[i].to;
        if(v != fa && !vis[v])  dfsdis(v, d + edges[i].dist, u);
    }
}

int calc(int u, int d)
{
    int ret = 0;
    num = 0;
    dfsdis(u, d, 0);
    sort(dis, dis + num);
    int i = 0, j = num - 1;
    while(i < j)
    {
        while(dis[i] + dis[j] > k && i < j) j--;
        ret += j - i;
        i++;
    }
    return ret;
}

void dfs(int u)
{
    root = u;
    dfssize(u, 0);
    dfsroot(u, u, 0);
    ans += calc(root, 0);
    vis[root] = 1;
    for(int i = edgehead[root]; i != -1; i = edges[i].next)
    {
        int v = edges[i].to;
        if(!vis[v])
        {
            ans -= calc(v, edges[i].dist);
            dfs(v);
        }
    }
}

void prework()
{
    memset(vis, 0, sizeof(vis));
    memset(edgehead, 0xff, sizeof(edgehead)); tot = 0;
    for(int i = 0; i < n - 1; ++i)
    {
        int a, b, c;
        scanf("%d %d %d", &a, &b, &c);
        addedge(a, b, c);
        addedge(b, a, c);
    }
}

void solve()
{
    ans = 0;
    dfs(1);
    printf("%d\n", ans);
}

int main()
{
    while(~scanf("%d %d", &n, &k))
    {
        if(!n && !k) break;
        prework();
        solve();
    }
    return 0;
}