bzoj3631: [JLOI2014]松鼠的新家(树上差分)
题目描述:松鼠的新家是一颗树,新家有n个房间,由n - 1条边连接。****要来参观,按一定的顺序参观n个房间,每到一个房间都要在那个房间拿走一个糖果(最后一个房间除外)。
问松鼠需要在每个房间各放几个糖果。
输入格式:第一行一个整数,表示房间的数量。
第二行n个整数,表示参观的顺序。
接下来n - 1行,每行两个整数表示树上相连的两个房间。
输入样例:
5
1 4 5 3 2
1 2
2 4
2 3
4 5
输出样例:
1
2
1
2
1
解析:可以说是树上差分的模板题,注意一下最后除了最开始参观的房间,其他的房间的糖果数都要减1.
代码如下:
1 #include<cstdio> 2 #include<vector> 3 #include<algorithm> 4 using namespace std; 5 6 const int maxn = 3e5 + 5; 7 int n, a[maxn], dep[maxn], fa[maxn][30], cnt[maxn]; 8 vector <int> ve[maxn]; 9 10 int read(void) { 11 char c; while (c = getchar(), c < '0' || c >'9'); int x = c - '0'; 12 while (c = getchar(), c >= '0' && c <= '9') x = x * 10 + c - '0'; return x; 13 } 14 15 void prepare(int u, int pre) { 16 dep[u] = dep[pre] + 1; 17 for (int i = 0; i < 29; ++ i) fa[u][i + 1] = fa[fa[u][i]][i]; 18 for (int i = 0; i < ve[u].size(); ++ i) { 19 int v = ve[u][i]; 20 if (v == pre) continue; 21 fa[v][0] = u; 22 prepare(v, u); 23 } 24 } 25 26 int getlca(int x, int y) { 27 if (x == 1 || y == 1) return 1; 28 if (dep[x] < dep[y]) swap(x, y); 29 for (int i = 29; i >= 0; -- i) 30 if (dep[fa[x][i]] >= dep[y]) { 31 x = fa[x][i]; 32 if (x == y) return x; 33 } 34 for (int i = 29; i >= 0; -- i) { 35 if (fa[x][i] != fa[y][i]) x = fa[x][i], y = fa[y][i]; 36 } 37 return fa[x][0]; 38 } 39 40 void dfs(int u, int pre) { 41 for (int i = 0; i < ve[u].size(); ++ i) { 42 int v = ve[u][i]; 43 if (v == pre) continue; 44 dfs(v, u); 45 cnt[u] += cnt[v]; 46 } 47 } 48 49 int main() { 50 n = read(); 51 for (int i = 1; i <= n; ++ i) a[i] = read(); 52 for (int i = 1; i < n; ++ i) { 53 int x = read(), y = read(); 54 ve[x].push_back(y); 55 ve[y].push_back(x); 56 } 57 prepare(1, 0); 58 for (int i = 1; i < n; ++ i) { //树上差分 59 cnt[a[i]] ++; cnt[a[i + 1]] ++; 60 int lca = getlca(a[i], a[i + 1]); 61 cnt[lca] --; cnt[fa[lca][0]] --; 62 } 63 dfs(1, 0); 64 cnt[a[1]] ++; //第一个点+1,在后面-1,抵消 65 for (int i = 1; i <= n; ++ i) printf("%d\n", cnt[i] - 1); //每个点都-1 66 return 0; 67 }