Count on a tree II SPOJ - COT2 树上莫队

#include<iostream>
#include<vector>
#include<algorithm>
#include<string>
#include<stack>
#include<cstring>
#include<set>
#include<iterator>
#include<list>
#include<deque>
#include<queue>
#include<map>
#include<cmath>
#include<sstream>
#include<cstdio>
#include<ctime>
#include<iomanip>
#include<unordered_map>
using namespace std;
typedef long long ll;
//typedef pair<int, int> P;
const int N = 4e4 + 10;
const int M = 1e5 + 10;
const int inf = 0x3f3f3f3f;
const ll INF = 0x3f3f3f3f3f3f3f3f;
int block_size;
struct Query
{
	int id, l, r;
	int ex;
	bool operator < (const Query &q)const
	{
		int b1 = (l - 1) / block_size;
		int b2 = (q.l - 1) / block_size;
		if (b1 != b2)
		{
			return b1 < b2;
		}
		else
		{
			return r < q.r;
		}
	}
}query[M];
int ans[M];
int rmq[2 * N];
struct ST
{
	int mm[2 * N];
	int dp[2 * N][20];
	void init(int n)
	{
		mm[0] = -1;
		for (int i = 1; i <= n; i++)
		{
			mm[i] = ((i&(i - 1)) == 0) ? mm[i - 1] + 1 : mm[i - 1];
			dp[i][0] = i;
		}
		for (int j = 1; j <= mm[n]; j++)
		{
			for (int i = 1; i + (1 << j) - 1 <= n; i++)
			{
				dp[i][j] = rmq[dp[i][j - 1]] < rmq[dp[i + (1 << (j - 1))][j - 1]] ? dp[i][j - 1] : dp[i + (1 << (j - 1))][j - 1];
			}
		}
	}
	int query(int a, int b)
	{
		if (a > b)
			swap(a, b);
		int k = mm[b - a + 1];
		return rmq[dp[a][k]] <= rmq[dp[b - (1 << k) + 1][k]] ? dp[a][k] : dp[b - (1 << k) + 1][k];
	}
}st;
int head[N], ver[2 * N], Next[2 * N], tot;
int F[2 * N], P[N], be[N], en[N], cnt;
int A[2 * N];
void init()
{
	tot = 0;
	memset(head, 0, sizeof(head));
}
void addedge(int u,int v)
{
	ver[++tot] = v, Next[tot] = head[u], head[u] = tot;
	ver[++tot] = u, Next[tot] = head[v], head[v] = tot;
}
int cur;
void dfs(int u, int pre, int dep)
{
	be[u] = ++cur;
	F[++cnt] = u;
	P[u] = cnt;
	rmq[cnt] = dep;
	for (int i = head[u]; i; i = Next[i])
	{
		int v = ver[i];
		if (v == pre)
			continue;
		dfs(v, u, dep + 1);
		F[++cnt] = u;
		rmq[cnt] = dep;
	}
	en[u] = ++cur;
}
void LCA_init(int root,int n)
{
	cnt = 0;
	dfs(root, root, 0);
	st.init(2 * n - 1);
}
int query_lca(int u, int v)
{
	return F[st.query(P[u], P[v])];
}
vector<int> vec;
int nowans = 0;
int vis[N], cou[N];
int a[N];
void change(int x)
{
	if (vis[x])
	{
		vis[x] ^= 1;
		cou[a[x]]--;
		if (cou[a[x]] == 0)
			nowans--;
	}
	else
	{
		vis[x] ^= 1;
		cou[a[x]]++;
		if (cou[a[x]] == 1)
		{
			nowans++;
		}
	}
}
int main()
{
	int n, m;
	scanf("%d%d", &n, &m);
	for (int i = 1; i <= n; i++)
	{
		scanf("%d", &a[i]);
		vec.push_back(a[i]);
	}
	sort(vec.begin(), vec.end());
	vec.erase(unique(vec.begin(), vec.end()), vec.end());
	for (int i = 1; i <= n; i++)
	{
		a[i] = lower_bound(vec.begin(), vec.end(), a[i]) - vec.begin() + 1;
	}
	for (int i = 1; i < n; i++)
	{
		int u, v;
		scanf("%d%d", &u, &v);
		addedge(u, v);
	}
	LCA_init(1, n);
	for (int i = 1; i <= n; i++)
	{
		A[be[i]] = i;
		A[en[i]] = i;
	}
	block_size = (int)sqrt(cur) + 1;
	for (int i = 1; i <= m; i++)
	{
		int u, v;
		scanf("%d%d", &u, &v);
		if (be[u] > be[v])
		{
			swap(u, v);
		}
		int lca = query_lca(u, v);
		if (lca == u)
		{
			query[i].id = i;
			query[i].ex = -1;
			query[i].l = be[u];
			query[i].r = be[v];
		}
		else
		{
			query[i].id = i;
			query[i].ex = be[lca];
			query[i].l = en[u];
			query[i].r = be[v];
		}
	}
	sort(query + 1, query + m + 1);
	int l = 1, r = 0;
	for (int i = 1; i <= m; i++)
	{
		while (r < query[i].r)
		{
			r++;
			change(A[r]);
		}
		while (r > query[i].r)
		{
			change(A[r]);
			r--;
		}
		while (l < query[i].l)
		{
			change(A[l]);
			l++;
		}
		while (l > query[i].l)
		{
			l--;
			change(A[l]);
		}
		if (query[i].ex != -1)
			change(A[query[i].ex]);
		ans[query[i].id] = nowans;
		if (query[i].ex != -1)
			change(A[query[i].ex]);
	}
	for (int i = 1; i <= m; i++)
	{
		printf("%d\n", ans[i]);
	}
	return 0;
}