poj 1741 Tree 点分治

Tree

Time Limit: 1000MS   Memory Limit: 30000K
Total Submissions: 30775   Accepted: 10301

Description

Give a tree with n vertices,each edge has a length(positive integer less than 1001). 
Define dist(u,v)=The min distance between node u and v. 
Give an integer k,for every pair (u,v) of vertices is called valid if and only if dist(u,v) not exceed k. 
Write a program that will count how many pairs which are valid for a given tree. 

Input

The input contains several test cases. The first line of each test case contains two integers n, k. (n<=10000) The following n-1 lines each contains three integers u,v,l, which means there is an edge between node u and v of length l. 
The last test case is followed by two zeros. 

Output

For each test case output the answer on a single line.

Sample Input

5 4
1 2 3
1 3 1
1 4 2
3 5 1
0 0

Sample Output

8

Source

LouTiancheng@POJ

题目大意:n个点的树,每条边有权重w,求有多少个点对(x, y)满足x到y的距离小于等于k

最基础的想法,对每个点都遍历一遍整个树,得到合法点对的个数,复杂度O(n2)无法满足需求。
注意到,以root为根的树上的路径,要么经过root,要么在root的子树内部,对于子树内部的路径可以通过递归得到解决;
对于经过root的路径,比如root有两个子树x和y,用depth[i]表示root到i节点的距离,那么合法的路径即为满足depth[i] + depth[j] <= k
(i∈x,j∈y)的个数,进一步的可以转换为 i∈(x, y),j∈(x, y)的合法点对 - i∈x,j∈x的合法点对 - i∈y,j∈y的合法点对
对于求depth[i] + depth[j] <= k 可以通过排序后双指针线性的得到答案,因此对于深度为D的树,每一层的复杂度为O(nlogn),所以总的复杂度为O(D*nlogn)
为了防止D退化为n,每次找root应为树的重心,求树的重心可以线性得到,因此总的复杂度为O(nlognlogn)

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

const int maxn = 1e4 + 10;
vector<pair<int, int> >G[maxn];
vector<int> vec;
int vis[maxn], sz[maxn], maxson[maxn], depth[maxn];
int n, k;
int ans, root, allnode;

void find_root(int u, int fa) {
    sz[u] = 1;
    maxson[u] = 0; //maxson[u]表示去除节点u之后节点数最多的树的节点数
    for (int i = 0; i < G[u].size(); i++) {
        int v = G[u][i].first;
        int w = G[u][i].second;
        if (v == fa || vis[v]) continue;
        find_root(v, u);
        maxson[u] = max(sz[v], maxson[u]);  //用子树v的节点数更新maxson[u]
        sz[u] += sz[v];
    }
    maxson[u] = max(maxson[u], allnode - sz[u]); 
    if (maxson[u] < maxson[root]) root = u;
}

//计算子树u下所有节点的depth
void dfs(int u, int fa) { 
    vec.push_back(depth[u]);
    sz[u] = 1;
    for (int i = 0; i < G[u].size(); i++) {
        int v = G[u][i].first;
        int w = G[u][i].second;
        if (v == fa || vis[v]) continue;
        depth[v] = depth[u] + w;
        dfs(v, u);
        sz[u] += sz[v];
    }
}

int calc(int u, int fa, int init) {
    vec.clear(); //节点u下的所有节点的depth会存在vec中
    depth[u] = init;
    dfs(u, fa);
    sort(vec.begin(), vec.end()); //对vec排序后用双指针扫一遍得到合法点对数
    int cnt = 0, l = 0, r = vec.size() - 1;
    for (; l < r;) {
        if (vec[l] + vec[r] <= k) {
            cnt += (r - l);
            l++;
        }
        else r--;
    }
    return cnt;
}

void solve(int u, int fa) {
    vis[u] = 1;
    ans += calc(u, fa, 0); //计算当前根下所有合法点对的个数
    for (int i = 0; i < G[u].size(); i++) {
        int v = G[u][i].first;
        int w = G[u][i].second;
        if (v == fa || vis[v]) continue;
        ans -= calc(v, u, w); //去除同一子树内部的合法点对
        root = 0;
        allnode = sz[v];
        find_root(v, u);
        solve(root, u);
    }
}

int main() {
    while (scanf("%d%d", &n, &k) != EOF) {
        if (n == 0 && k == 0) break;
        for (int i = 1; i <= n; i++) G[i].clear();
        ans = 0;
        maxson[0] = n + 1;
        memset(vis, 0, sizeof(vis));
        for (int i = 1; i < n; i++) {
            int u, v, w;
            scanf("%d%d%d", &u, &v, &w);
            G[u].push_back({v, w});
            G[v].push_back({u, w});
        }
        root = 0;
        allnode = n;
        find_root(1, -1);
        solve(root, -1);
        printf("%d\n", ans);
    }
    return 0;
}

猜你喜欢

转载自blog.csdn.net/KIDGIN7439/article/details/83507819
今日推荐