点分治【模板】

推荐学习博客:
https://blog.csdn.net/a_forever_dream/article/details/81778649
https://www.cnblogs.com/bztMinamoto/p/9489473.html
https://www.luogu.org/blog/user9012/dian-fen-zhi-lve-xie

贴个模板(洛谷Tree):

///#include<bits/stdc++.h>
///#include<unordered_map>
///#include<unordered_set>
#include<iostream>
#include<algorithm>
#include<cstdio>
#include<cstdlib>
#include<cstring>
#include<string>
#include<cmath>
#include<queue>
#include<bitset>
#include<set>
#include<stack>
#include<map>
#include<list>
#include<new>
#include<vector>

#define MT(a, b) memset(a,b,sizeof(a)
#define lowbit(x) (x&(-x))
using namespace std;
typedef long long ll;
const double pai = acos(-1.0);
const double E = 2.718281828459;
const ll mod = 1e9 + 7;
const int INF = 0x3f3f3f3f;
const double esp = 1e-6;
const int maxn = 4e4 + 5;

int n, m, k;
int root, MAX, size, sz[maxn];
/// root子树的重心,size当前分治时子树的大小
///sz[i]记录i子树的儿子数量
int vis[maxn];  ///标记点是否已经分治过了
int dis[maxn], cnt; ///子树到点的距离,cnt记录数量
ll ans = 0;     ///ans=可能的答案-不可能的答案
struct node {
    int e, c, p;
} load[maxn << 1];
int head[maxn], sign;

void add_edge(int s, int e, int c) {
    load[++sign] = node{e, c, head[s]};
    head[s] = sign;
}

void get_root(int s, int pre) {
    sz[s] = 1;
    int m = 0;
    for (int i = head[s], e; ~i; i = load[i].p) {
        e = load[i].e;
        if (e != pre && !vis[e]) {
            get_root(e, s);
            sz[s] += sz[e];
            m = max(m, sz[e]);
        }
    }
    m = max(m, size - sz[s]);
    if (MAX > m)
        MAX = m, root = s;
}

void get_dis(int s, int pre, int depth) {
    dis[++cnt] = depth;
    for (int i = head[s], e; ~i; i = load[i].p) {
        e = load[i].e;
        if (e != pre && !vis[e]) {
            get_dis(e, s, depth + load[i].c);
        }
    }
}

ll get_ans(int s, int depth) {
    cnt = 0;
    get_dis(s, -1, depth);
    sort(dis + 1, dis + 1 + cnt);
    ll all = 0;
    for (int i = 1; i <= cnt; i++) {
        int l = i + 1, r = cnt, mid, temp = -1;
        while (l <= r) {            ///二分查找答案
            mid = (l + r) >> 1;
            if (dis[mid] + dis[i] <= k) {
                temp = mid;
                l = mid + 1;
            } else
                r = mid - 1;
        }
        if (temp != -1)
            all += temp - i;
    }
    return all;
}

void divide(int s,int now_sz) {
    vis[s] = 1;
    ans += get_ans(s, 0);   ///加上可能的答案
    for (int i = head[s], e; ~i; i = load[i].p) {
        e = load[i].e;
        if (vis[e])
            continue;
        ans -= get_ans(e, load[i].c);   ///减去子树中不可能的答案
        size = sz[e] > sz[s] ? now_sz- sz[s] : sz[e];    ///判断节点e和s的关系
        MAX = INF;        ///更新子树的信息
        get_root(e, -1);
        divide(root,size);                   ///继续分治
    }
}
void init() {
    size = n, sign = 0, MAX = INF;
    memset(vis, 0, sizeof(vis));
    memset(head, -1, sizeof(head));
}

int main() {
    int s, e, c;
    scanf("%d", &n);
    init();
    for (int i = 1; i < n; i++) {
        scanf("%d %d %d", &s, &e, &c);
        add_edge(s, e, c);
        add_edge(e, s, c);
    }
    scanf("%d", &k);
    get_root(1, -1);
    divide(root,n);
    printf("%lld\n", ans);
    return 0;
}

 

全部评论

相关推荐

点赞 收藏 评论
分享
牛客网
牛客企业服务