AcWing 252. 树(点分治模版题)

题目链接

题目大意

  略

解法1：点分治+归并排序+二分

  对于当前的重心，首先对其一棵子树求距离，然后二分之前的子树中符合条件的距离的数量，然后用归并来降低排序的时间复杂度。

代码

const int maxn = 1e4+10;
int n, m, rt, tot, tota, totb, totc; 
bool vis[maxn]; ll ans;
int d[maxn], sz[maxn], a[maxn], b[maxn], c[maxn], h[maxn], mx[maxn];
struct E {
    int to, w, nxt;
} e[maxn<<1];
void add(int u, int v, int w) {
    e[++tot] = {v, w, h[u]};
    h[u] = tot;
}
void init() {
    ans = tot = 0;
    for (int i = 0;  i<=n; ++i) h[i] = 0, vis[i] = 0;
}
void getrt(int u, int p, int szr) { 
    sz[u] = 1; mx[u] = 0;
    for (int i = h[u]; i; i=e[i].nxt) {
        int v = e[i].to;
        if (v==p || vis[v]) continue;
        getrt(v, u, szr);
        sz[u] += sz[v];
        if (sz[v]>mx[u]) mx[u] = sz[v];
    }
    if (szr-sz[u]>mx[u]) mx[u] = szr-sz[u];
    if (!rt || mx[u]<mx[rt]) rt = u;
}
void getdis(int u, int p) {
    for (int i = h[u]; i; i=e[i].nxt) {
        int v = e[i].to, w = e[i].w;
        if (v==p || vis[v]) continue;
        d[v] = d[u]+w;
        b[++totb] = d[v];
        getdis(v, u);
    }
}
void merge() {
    int t1 = 1, t2 = 1; totc = 1;
    while(t1<=tota || t2<=totb) {
        if (t2>totb || (t1<=tota && a[t1]<b[t2])) c[totc++] = a[t1++];
        else c[totc++] = b[t2++];
    }
    for (int i = 1; i<=tota+totb; ++i) a[i] = c[i];
    tota = tota+totb;
}
void calc(int u) {
    d[u] = tota = 0;
    for (int i = h[u]; i; i=e[i].nxt) {
        int v = e[i].to, w = e[i].w;
        if (vis[v]) continue;
        d[v] = d[u]+w; 
        totb = 0;
        b[++totb] = d[v]; //存同一棵子树中的所有值
        getdis(v, u);
        sort(b+1, b+totb+1);
        for (int i = 1; i<=totb; ++i) {
            if (b[i]>m) continue;
            else {
                if (b[i]<=m) ++ans;
                int r = upper_bound(a+1, a+tota+1, m-b[i])-a; //在之前的子树中找值
                ans += r-1;
            }
        }
        merge();
    }
}
void div(int u) {
    vis[u] = 1;
    calc(u); //计算贡献
    for (int i = h[u]; i; i=e[i].nxt) {
        int v = e[i].to;
        if (vis[v]) continue;
        rt = 0; sz[rt] = INF;
        int t = sz[v];
        getrt(v, -1, t); //第一遍求重心
        getrt(rt, -1, t); //第二遍求以重心为根的子树大小
        div(rt);
    }
}
int main(void) {
    while(~scanf("%d%d", &n, &m) && (n||m)) {
        init();
        for (int i = 1, a, b, c; i<n; ++i) {
            scanf("%d%d%d", &a, &b, &c);
            add(a, b, c);
            add(b, a, c);
        }
        rt = 0, sz[rt] = INF;
        mx[rt] = INF;
        getrt(1, -1, n); //第一遍求重心
        getrt(rt, -1, n); //第二遍求以重心为根的子树大小
        div(rt);
        printf("%lld
", ans);
    }
    return 0;
}  

解法2：点分治+树状数组

  思路与解法1类似，只不过把之前的子树的权值都加到了树状数组上

代码

const int maxn = 1e4+10;
const int maxm = 1e7+10;
int n, m, rt, tot, q1, q2; 
bool vis[maxn]; ll ans;
int c[maxm];
int d[maxn], sz[maxn], h[maxn], mx[maxn], que1[maxn], que2[maxn];
void add2(int x, int y) {
    ++x; //题目数据有权值为1的边，整体加1
    while(x<maxm) {
        c[x] += y;
        x += x&-x;
    }
}
int ask(int x) {
    ++x; //题目数据有权值为1的边，整体加1
    int sum = 0;
    while(x) {
        sum += c[x];
        x -= x&-x;
    }
    return sum;
}
struct E {
    int to, w, nxt;
} e[maxn<<1];
void add(int u, int v, int w) {
    e[++tot] = {v, w, h[u]};
    h[u] = tot;
}
void init() {
    ans = tot = 0;
    for (int i = 0;  i<=n; ++i) h[i] = 0, vis[i] = 0;
}
void getrt(int u, int p, int szr) { 
    sz[u] = 1; mx[u] = 0;
    for (int i = h[u]; i; i=e[i].nxt) {
        int v = e[i].to;
        if (v==p || vis[v]) continue;
        getrt(v, u, szr);
        sz[u] += sz[v];
        if (sz[v]>mx[u]) mx[u] = sz[v];
    }
    if (szr-sz[u]>mx[u]) mx[u] = szr-sz[u];
    if (!rt || mx[u]<mx[rt]) rt = u;
}
void getdis(int u, int p) {
    for (int i = h[u]; i; i=e[i].nxt) {
        int v = e[i].to, w = e[i].w;
        if (v==p || vis[v]) continue;
        d[v] = d[u]+w;
        que1[++q1] = que2[++q2] = d[v];
        getdis(v, u);
    }
}
void calc(int u) {
    d[u] = q1 = 0;
    for (int i = h[u]; i; i=e[i].nxt) {
        int v = e[i].to, w = e[i].w;
        if (vis[v]) continue;
        d[v] = d[u]+w; 
        q2 = 0;
        que1[++q1] = que2[++q2] = d[v];
        getdis(v, u);
        for (int j = 1; j<=q2; ++j) {
            if (que2[j]>m) continue;
            ans += ask(m-que2[j])+1;
        }
        for (int j = 1; j<=q2; ++j) add2(que2[j], 1);
    }
    for (int i = 1; i<=q1; ++i) add2(que1[i], -1);
}
void div(int u) {
    vis[u] = 1;
    calc(u); //计算贡献
    for (int i = h[u]; i; i=e[i].nxt) {
        int v = e[i].to;
        if (vis[v]) continue;
        rt = 0; mx[rt] = INF;
        int t = sz[v];
        getrt(v, -1, t); //第一遍求重心
        getrt(rt, -1, t); //第二遍求以重心为根的子树大小
        div(rt);
    }
}
int main(void) {
    while(~scanf("%d%d", &n, &m) && (n||m)) {
        init();
        for (int i = 1, a, b, c; i<n; ++i) {
            scanf("%d%d%d", &a, &b, &c);
            ++a, ++b;
            add(a, b, c);
            add(b, a, c);
        }
        rt = 0, sz[rt] = INF;
        mx[rt] = INF;
        getrt(1, -1, n); //第一遍求重心
        getrt(rt, -1, n); //第二遍求以重心为根的子树大小
        div(rt);
        printf("%lld
", ans);
    }
    return 0;
}  

解法3：点分治+尺取

  将以当前重心为根的树中所有的点到重心的距离排序，然后尺取，可以当左指针右移时，右指针只会向左移动，此时r-l再减去和左指针指向的点在同一棵子树中的点的数量就是答案。

代码

const int maxn = 1e4+10;
int n, m, rt, tot, tota; ll ans;
bool vis[maxn];
int d[maxn], sz[maxn], a[maxn], b[maxn], h[maxn], cnt[maxn], mx[maxn];
struct E {
    int to, w, nxt;
} e[maxn<<1];
void add(int u, int v, int w) {
    e[++tot] = {v, w, h[u]};
    h[u] = tot;
}
void init() {
    ans = tot = 0;
    for (int i = 0;  i<=n; ++i) h[i] = 0, vis[i] = 0;
}
void getrt(int u, int p, int szr) { //szr存当前的子树中的节点个数
    sz[u] = 1; mx[u] = 0;
    for (int i = h[u]; i; i=e[i].nxt) {
        int v = e[i].to;
        if (v==p || vis[v]) continue;
        getrt(v, u, szr);
        sz[u] += sz[v];
        if (sz[v]>mx[u]) mx[u] = sz[v];
    }
    if (szr-sz[u]>mx[u]) mx[u] = szr-sz[u];
    if (!rt || mx[u]<mx[rt]) rt = u;
}
void getdis(int u, int p, int x) {
    b[u] = x, ++cnt[x], a[++tota] = u;
    for (int i = h[u]; i; i=e[i].nxt) {
        int v = e[i].to, w = e[i].w;
        if (v==p || vis[v]) continue;
        d[v] = d[u]+w;
        getdis(v, u, x);
    }
}
bool cmp(int a, int b) {
    return d[a]<d[b];
}
void calc(int u) {
    d[u] = tota = 0;
    b[u] = u, ++cnt[u], a[++tota] = u;  //存下每个点为根的子树大小
    for (int i = h[u]; i; i=e[i].nxt) {
        int v = e[i].to, w = e[i].w;
        if (vis[v]) continue;
        d[v] = d[u]+w; 
        getdis(v, u, v); 
    }
    //cout << tota << endl;
    sort(a+1, a+tota+1, cmp);
    int l = 1, r = tota; --cnt[b[a[l]]];
    //cout << l << ' ' << r << endl;
    while(l<r) {
        while(r>l && d[a[l]]+d[a[r]]>m) --cnt[b[a[r]]], --r; //cnt存储l+1 ~ r之间同一子树的点数
        if (l>=r) break;
        ans += r-l-cnt[b[a[l]]]; //计算时减去与a[l]同一子树的贡献
        //cout << ans << endl;
        ++l;
        --cnt[b[a[l]]];
    }
    for (int i = 0; i<=tota; ++i) b[a[i]] = cnt[a[i]] = 0;
}
void div(int u) {
    vis[u] = 1;
    calc(u); //计算贡献
    for (int i = h[u]; i; i=e[i].nxt) {
        int v = e[i].to;
        if (vis[v]) continue;
        rt = 0; mx[rt] = INF;
        int t = sz[v];
        getrt(v, -1, t);
        getrt(rt, -1, t);
        div(rt);
    }
}
int main(void) {
    while(~scanf("%d%d", &n, &m) && (n||m)) {
        init();
        for (int i = 1, a, b, c; i<n; ++i) {
            scanf("%d%d%d", &a, &b, &c);
            ++a, ++b;
            add(a, b, c);
            add(b, a, c);
        }
        rt = 0, mx[rt] = INF;
        getrt(1, -1, n);
        getrt(rt, -1, n); //第二遍求重心保证sz是以重心为根算出来的
        div(rt);
        printf("%lld
", ans);
    }
    return 0;
}