BZOJ3653——谈笑风生(dfs序 && 树状数组)

现在有一颗以 1 为根的有根树，有 m 个询问，每次给定a, k, 求满足条件的三元组(a, b, c)(a 就是已经给定的那个)，条件：1、a, b距离不超过 k，2、a，b都是c的祖先。

首先比较显然的是需要分类：

a). b 在 a 的上面所构成的三元组，显然这样的可以直接统计出来，即 min(dep[a], k) * (size[a] - 1)，每一个上面的和 a 的子树中每个点都构成了一个三元组。

b). b 在 a 的子树中。因为我们直接限定了范围为 a 的子树，而且与 a 的距离不超过 k，所以就明晰了，把询问离线，按需要询问的深度排序，这样不超过 k 这个条件就可以搞定了，然后dfs序用最常用的树状数组维护就可以了。

#include <cstdio>
#include <cstring>
#include <algorithm>

using namespace std;

const int MAX_N = 300005;
typedef long long ll;

struct node {
	int v, next;
} E[MAX_N << 1];
int head[MAX_N], top = 0;

void add(int u, int v)
{
	E[++ top].v = v; E[top].next = head[u]; head[u] = top;
}

int n, m;
int sz[MAX_N], fa[MAX_N], dep[MAX_N], tm[MAX_N], st[MAX_N], tsz = 0;

void dfs(int x, int last)
{
	sz[x] = 1; fa[x] = last;
	tm[st[x] = ++ tsz] = x;
	for (int i = head[x]; i; i = E[i].next) {
		if (E[i].v == last) continue;
		dep[E[i].v] = dep[x] + 1;
		dfs(E[i].v, x);
		sz[x] += sz[E[i].v];
	}
}

void init()
{
	scanf("%d%d", &n, &m);
	for (int i = 1, u, v; i < n; i ++) {
		scanf("%d%d", &u, &v);
		add(u, v); add(v, u);
	}
}

int q[MAX_N];
struct qry {
	int dep, id, x;
} qr[MAX_N];

inline bool cmp(qry a, qry b)
{
	return a.dep < b.dep;
}

ll c[MAX_N << 1];
ll ans[MAX_N];

void Add(int x, ll w) { for (; x <= n; x += x & -x) c[x] += w;}
ll Qry(int x) { ll ret = 0; for (; x > 0; x -= x & -x) ret += c[x]; return ret;}

void doit()
{
	dfs(1, 0);
	int hd = 0, tl = 0;
	q[++ tl] = 1;
	while (hd < tl) {
		int x = q[++ hd];
		for (int i = head[x]; i; i = E[i].next) 
			if (E[i].v != fa[x]) q[++ tl] = E[i].v;
	}
	for (int i = 1; i <= m; i ++) {
		int x, k; scanf("%d%d", &x, &k);
		qr[i].x = x; qr[i].id = i; qr[i].dep = dep[x] + k;
		int tmp = min(k, dep[x]);
		ans[i] = 1ll * tmp * (sz[x] - 1);
	}
	sort(qr + 1, qr + m + 1, cmp);
	int tp = 0;
	for (int i = 1; i <= m; i ++) {
		while (tp < n && dep[q[tp + 1]] <= qr[i].dep) tp ++, Add(st[q[tp]], (ll)(sz[q[tp]] - 1));
		ans[qr[i].id] += Qry(st[qr[i].x] + sz[qr[i].x] - 1) - Qry(st[qr[i].x]);
	}
	for (int i = 1; i <= m; i ++) printf("%lld\n", ans[i]);
}

int main()
{
	init();
	doit();
	return 0;
}