ACM模板库

ACM模板库


点分治

下面求树上小于等于k的路径数量(路径两端点不同)

#include <iostream>
#include <vector>
#include <algorithm>
using namespace std;
const int maxn = 1e5+5;
typedef long long ll;
typedef pair<ll, int> Edge;
vector<Edge> mp[maxn];
int n, k;
//siz表示子树大小,maxson表示重儿子子树的大小 
int siz[maxn], maxson[maxn] = {maxn}, root;
bool vis[maxn];
vector<ll> dis;
int subn;//求根时子树大小 
ll ans = 0;
void getroot(int u, int fa){//n为树的大小,求重心 
    siz[u] = 1; 
    maxson[u] = 0;
    for(int i = 0; i < mp[u].size(); i++){
        int v = mp[u][i].second;
        if(vis[v] || v == fa)   continue;
        getroot(v, u);
        siz[u] = siz[u] + siz[v];
        maxson[u] = max(maxson[u], siz[v]);
    }
    maxson[u] = max(maxson[u], subn - siz[u]);
    if(maxson[u] < maxson[root]) root = u;
}
void getdis(int u, int fa, ll dist){
    dis.push_back(dist);
    for(int i = 0; i < mp[u].size(); i++){
        int v = mp[u][i].second;
        if(vis[v] || v == fa)   continue;
        getdis(v, u, dist + mp[u][i].first);
    }
    return;
}
int calc(int u, ll len){ //这个用来求解的函数,自定义 
    int cnt = 0;
    ll ans = 0;
    dis.clear();
    getdis(u, 0, len);
    sort(dis.begin(), dis.end());   
    int l = 0, r = dis.size()-1;
    while(l <= r){
        if(dis[l] + dis[r] <= k){
            ans = ans + r - l; 
            l++;
        }
        else r--;
    }
    return ans;
}
void divide(int u){
    ans = ans + calc(u, 0);
    vis[u] = true;
    for(int i = 0; i < mp[u].size(); i++){
        int v = mp[u][i].second;
        if(vis[v])  continue;
        ans = ans - calc(v, mp[u][i].first);
        subn = siz[v];
        root = 0;
        getroot(v, 0);
        divide(root);
    }
    return;
}
void init(){
    subn = n;
    ans = 0;
    for(int i = 1; i <= n; i++){
        vis[i] = false; 
        mp[i].clear();
    }
}
int main(){
    while(scanf("%d%d", &n, &k)){
        if(n == 0 && k == 0)    return 0;
        init();
        for(int i = 1; i < n; i++){
            int u, v;
            ll w;
            scanf("%d%d%lld", &u, &v, &w);
            mp[u].push_back(Edge(w, v));
            mp[v].push_back(Edge(w, u));
        }
        getroot(1, 0);
        divide(root);
        cout << ans << endl;
    }

} 

页面列表

ITEM_HTML