P4178 Tree(点分治)

题目描述
给定一棵 nn 个节点的树,每条边有边权,求出树上两点距离小于等于 kk 的点对数量

思路:
就是裸的点分治,每次找到重心遍历,然后求出所有路径双指针搞一下。

数据开到了4e4,卡常,极限优化常数后,最后发现是unordered_map的问题,说好的O(1)查询呢??

#include <cstdio>
#include <cstring>
#include <algorithm>
#include <cmath>
#include <vector>
#include <queue>
#include <iostream>
#include <map>
#include <string>

using namespace std;

typedef long long ll;

const int maxn = 4e4 + 7;

int head[maxn],nex[maxn * 2],to[maxn * 2],val[maxn * 2],tot;
int a[maxn],b[maxn],d[maxn],cnt[maxn],vis[maxn],w[maxn],siz[maxn];
int ans,n,k,num,pos,ans_siz;

void add(int x,int y,int z) {
    to[++tot] = y;
    val[tot] = z;
    nex[tot] = head[x];
    head[x] = tot;
}

void dfs_find(int s,int x,int fa) {
    vis[x] = 1;
    siz[x] = 1;
    int max_part = 0;
    for(int i = head[x];i;i = nex[i]) {
        int y = to[i];
        if(y == fa || w[y]) continue;
        dfs_find(s,y,x);
        siz[x] += siz[y];
        max_part = max(max_part,siz[y]);
    }
    max_part = max(max_part,s - siz[x]);
    if(max_part < ans_siz) {
        ans_siz = max_part;
        pos = x;
    }
}

void dfs(int x,int fa) {
    vis[x] = 1;
    for(int i = head[x];i;i = nex[i]) {
        int y = to[i],z = val[i];
        if(y == fa || w[y]) continue;
        ++cnt[b[x]];a[++num] = y;b[y] = b[x];
        d[y] = d[x] + z;
        dfs(y,x);
    }
}

int cmp(int x,int y) {
    return d[x] < d[y];
}

void work(int s,int x) { //总大小s,根节点x
    ans_siz = s;
    dfs_find(s,x,-1); //找到重心
    num = 1;
    a[num] = b[pos] = pos;
    ++cnt[pos];
    w[pos] = 1;
    
    for(int i = head[pos];i;i = nex[i]) { //当前根节点(重心)是pos
        int y = to[i],z = val[i];
        if(w[y]) continue;
        ++cnt[y];a[++num] = b[y] = y;
        d[y] = z;
        dfs(y,pos);
    }
    
    sort(a + 1,a + 1 + num,cmp);
    int l = 1,r = num;
    --cnt[b[a[1]]];
    while(l < r) {
        while(d[a[l]] + d[a[r]] > k) {
            --cnt[b[a[r]]];r--;
        }
        ans += r - l - cnt[b[a[l]]];
        ++l;--cnt[b[a[l]]];
    }
    
    for(int i = 1;i <= num;i++) {
        cnt[b[a[i]]] = 0;
        d[a[i]] = 0;
    }
    int now = pos;
    dfs_find(s,now,-1); //更新siz函数
    for(int i = head[now];i;i = nex[i]) {
        int y = to[i];
        if(w[y]) continue;
        work(siz[y],y);
    }
}

void solve() {
    tot = 0;
    memset(head,0,sizeof(head));
    memset(w,0,sizeof(w));
    ans = 0;
    for(int i = 1;i < n;i++) {
        int x,y,z;scanf("%d%d%d",&x,&y,&z);
        add(x,y,z);add(y,x,z);
    }
    scanf("%d",&k);
    work(n,1);
    printf("%d\n",ans);
}

int main() {
    while(~scanf("%d",&n) && n) {
        solve();
    }
    return 0;
}

©️2020 CSDN 皮肤主题: 黑客帝国 设计师:上身试试 返回首页