HDU-3974-Assign the task

HDU-3974-Assign the task

  题意是给定一棵树，一种操作是指定一个点，这个点及这个点的的子树被染色，另一种操作是指定一个点，问这个点的颜色。可以通过dfs序将这棵树放在线段树上，定义两个数组dfn[n]和Count[n]，dfn保存的是每个节点的dfs序，Count[u]保存的是以u为根结点子树的节点总数，所以每次将节点u及u的子树染色就是染色线段树上的区间[dfn[u],dfn[u]+Count[u]-1]染色，询问某个节点u的颜色就是询问线段树上dfn[u]的值。一次搜索可以将这两个数组维护出来。

#include<iostream>
#include<algorithm>
#include<cstdio>
#include<cstring>
using namespace std;
const int MAX_N = 5e4 + 7;
struct edge {
	int v, nxt;
}e[MAX_N];
int n, m;
struct node {
	int l, r, v;
}tr[MAX_N * 4];
int w[MAX_N];
int dfn[MAX_N], Count[MAX_N];
int head[MAX_N], cnt = 1;
int tag = 0;
void add(int u, int v) {
	e[cnt].v = v;
	e[cnt].nxt = head[u];
	head[u] = cnt++;
}
int dfs(int u) {
	dfn[u] = ++tag;
	int temp = 1;
	for (int i = head[u]; ~i; i = e[i].nxt) {
		int v = e[i].v;
		if (!dfn[v]) {
			temp += dfs(v);
		}
	}
	Count[u] = temp;
	return temp;
}
void pushdown(int u) {
	if (tr[u].v == -1)return;
	tr[u<<1].v = tr[u].v;
	tr[u<<1|1].v = tr[u].v;
	tr[u].v = -1;
}

void build(int u, int l, int r) {//其实这道题不需要维护左右区间，但是出于个人习惯还是写了建树的函数
	if (l == r)tr[u] = { l,r,w[r] };
	else {
		tr[u] = { l,r,-1 };
		int mid = l + r >> 1;
		build(u << 1, l, mid);
		build(u << 1 | 1, mid + 1, r);
	}
}

int query(int u, int p) {
	if (tr[u].l == tr[u].r)return tr[u].v;
	pushdown(u);
	int mid = tr[u].l + tr[u].r >> 1;
	if (p <= mid)return query(u << 1, p);
	else return query(u << 1 | 1, p);
}
void modify(int u, int l, int r, int t) {
	if (tr[u].l >= l && tr[u].r <= r) {
		tr[u].v = t;
		return;
	}
	pushdown(u);
	int mid = tr[u].l + tr[u].r >> 1;
	if (l <= mid)modify(u << 1, l, r, t);
	if (r >= mid + 1)modify(u << 1 | 1, l, r, t);
}

int main() {
	int t;
	cin >> t;
	for (int k = 1; k <= t; k++) {
		cin >> n;
		memset(head, -1, sizeof head);
		memset(dfn, 0, sizeof dfn);
		memset(Count, 0, sizeof Count);
		cnt = 1, tag = 0;
		int u, v;
		for (int i = 1; i <= n - 1; i++) {
			scanf("%d%d", &u, &v);
			add(v, u);
			Count[u]++;
		}
		for (int i = 1; i <= n; i++) {
			if (!Count[i]) {
				dfs(i);
				break;
			}
		}
		scanf("%d", &m);
		memset(w, -1, sizeof w);
		build(1, 1, n);
		char op[5];
		printf("Case #%d:\n", k);
		for (int i = 1; i <= m; i++) {
			scanf("%s", op);
			if (op[0] == 'C') {
				scanf("%d", &u);
				printf("%d\n", query(1, dfn[u]));
			}
			else {
				scanf("%d%d", &u, &v);
				modify(1, dfn[u], dfn[u] + Count[u] - 1, v);
			}
		}
	}
	return 0;
}