Tree

题目描述

Give a tree with $n$ vertices, each edge has a length (positive integer less than $1001$).

Define $dist(u, v)= \text{The min distance between node} \ u \ \text{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.

题意概述

给定一棵有$n$个节点的树以及每条边的权值，问有多少个点对之间的距离不大于$k$。

数据范围：$2 \le n \le 10000$。

算法分析

考虑对于一个节点$x$以及它的子树，其中距离不大于$k$的点对要么在$x$的同一棵子树中，要么在$x$的不同子树中。后者可以递归解决，下面来解决前者。

可以统计出子树中每个节点到$x$的距离$dist_{u, x}$。显然有

$$
dist_{u, v} \le k \text{且路径经过} x \Leftrightarrow dist_{u, x} + dist_{x, v} \le k \text{且} u \text{和} v \text{在不同子树中}
$$

枚举$x$的子树，将前$i$棵子树中所有节点到$x$的距离用 Treap 维护，在处理第$(i+1)$棵子树的节点$t$时把答案加上 Treap 中小于等于$k-dist_{t, x}$的数的个数，即可得到答案。

代码

#include <cstdio>
#include <cstdlib>
#include <cstring>
#include <algorithm>
using namespace std;
struct edge {
    int v, w, nxt;
} e[20001];
struct treap {
    int tot, root;
    struct node_type {
        int val, cnt, size, rank, child[2];
    } a[10001];
    void init() { tot = root = 0; }
    void update(int t) { a[t].size = a[a[t].child[0]].size + a[a[t].child[1]].size + a[t].cnt; }
    void turn(int &t, int dire) {
        int p = a[t].child[!dire];
        a[t].child[!dire] = a[p].child[dire], a[p].child[dire] = t;
        update(t), update(p), t = p;
    }
    void insert(int &t, int val) {
        if (!t) {
            t = ++tot, a[t].rank = rand(), a[t].cnt = a[t].size = 1, a[t].val = val;
            a[t].child[0] = a[t].child[1] = 0;
            return;
        }
        ++a[t].size;
        if (val == a[t].val) ++a[t].cnt;
        else if (val < a[t].val) {
            insert(a[t].child[0], val);
            if (a[a[t].child[0]].rank < a[t].rank) turn(t, 1);
        } else {
            insert(a[t].child[1], val);
            if (a[a[t].child[1]].rank < a[t].rank) turn(t, 0);
        }
    }
    int query(int t, int val) {
        if (!t) return 0;
        if (val == a[t].val) return a[a[t].child[0]].size;
        else if (val < a[t].val) return query(a[t].child[0], val);
        else return a[a[t].child[0]].size + a[t].cnt + query(a[t].child[1], val);
    }
} tree;
long long n, k, nume, root, tot, ans, h[10001], size[10001], f[10001], val[10001];
bool vis[10001];
void add_edge(int u, int v, int w) {
    e[++nume].v = v, e[nume].w = w, e[nume].nxt = h[u], h[u] = nume;
    e[++nume].v = u, e[nume].w = w, e[nume].nxt = h[v], h[v] = nume;
}
void get_root(int t, int fa) {
    size[t] = 1, f[t] = 0;
    for (int i = h[t]; i; i = e[i].nxt) {
        if (!vis[e[i].v] && e[i].v != fa) {
            get_root(e[i].v, t);
            size[t] += size[e[i].v];
            f[t] = max(f[t], size[e[i].v]);
        }
    }
    f[t] = max(f[t], tot - size[t]);
    if (f[t] < f[root]) root = t;
}
void get_dist(int t, int fa, int flag) {
    if (!flag) tree.insert(tree.root, val[t]);
    else ans += tree.query(tree.root, k - val[t] + 1);
    for (int i = h[t]; i; i = e[i].nxt) {
        if (!vis[e[i].v] && e[i].v != fa) {
            val[e[i].v] = val[t] + e[i].w;
            get_dist(e[i].v, t, flag);
        }
    }
}
void solve(int t) {
    vis[t] = true;
    tree.init(), tree.insert(tree.root, 0);
    for (int i = h[t]; i; i = e[i].nxt) {
        if (!vis[e[i].v]) {
            val[e[i].v] = e[i].w;
            get_dist(e[i].v, t, 1);
            get_dist(e[i].v, t, 0);
        }
    }
    for (int i = h[t]; i; i = e[i].nxt) {
        if (!vis[e[i].v]) {
            root = 0, tot = size[e[i].v];
            get_root(e[i].v, t);
            solve(root);
        }
    }
}
int main()
{
    while (scanf("%lld%lld", &n, &k)) {
        if (!n) break;
        memset(vis, 0, sizeof(vis));
        memset(h, 0, sizeof(h));
        ans = nume = 0;
        for (int i = 1; i < n; ++i) {
            int u, v, w;
            scanf("%d%d%d", &u, &v, &w);
            add_edge(u ,v, w);
        }
        tot = f[0] = n, root = 0;
        get_root(1, 0);
        solve(root);
        printf("%lld\n", ans);
    }
    return 0;
}