CSU1811: Tree Intersection

  1 #include <bits/stdc++.h>
  2 using namespace std;
  3 namespace my_header {
  4 #define pb push_back
  5 #define mp make_pair
  6 #define pir pair<int, int>
  7 #define vec vector<int>
  8 #define pc putchar
  9 #define clr(t) memset(t, 0, sizeof t)
 10 #define pse(t, v) memset(t, v, sizeof t)
 11 #define bl puts("")
 12 #define wn(x) wr(x), bl
 13 #define ws(x) wr(x), pc(‘ ‘)
 14     const int INF = 0x3f3f3f3f;
 15     typedef long long LL;
 16     typedef double DB;
 17     inline char gchar() {
 18         char ret = getchar();
 19         for(; (ret == ‘\n‘ || ret == ‘\r‘ || ret == ‘ ‘) && ret != EOF; ret = getchar());
 20         return ret; }
 21     template<class T> inline void fr(T &ret, char c = ‘ ‘, int flg = 1) {
 22         for(c = getchar(); (c < ‘0‘ || ‘9‘ < c) && c != ‘-‘; c = getchar());
 23         if (c == ‘-‘) { flg = -1; c = getchar(); }
 24         for(ret = 0; ‘0‘ <= c && c <= ‘9‘; c = getchar())
 25             ret = ret * 10 + c - ‘0‘;
 26         ret = ret * flg; }
 27     inline int fr() { int t; fr(t); return t; }
 28     template<class T> inline void fr(T&a, T&b) { fr(a), fr(b); }
 29     template<class T> inline void fr(T&a, T&b, T&c) { fr(a), fr(b), fr(c); }
 30     template<class T> inline char wr(T a, int b = 10, bool p = 1) {
 31         return a < 0 ? pc(‘-‘), wr(-a, b, 0) : (a == 0 ? (p ? pc(‘0‘) : p) : 
 32             (wr(a/b, b, 0), pc(‘0‘ + a % b)));
 33     }
 34     template<class T> inline void wt(T a) { wn(a); }
 35     template<class T> inline void wt(T a, T b) { ws(a), wn(b); }
 36     template<class T> inline void wt(T a, T b, T c) { ws(a), ws(b), wn(c); }
 37     template<class T> inline void wt(T a, T b, T c, T d) { ws(a), ws(b), ws(c), wn(d); }
 38     template<class T> inline T gcd(T a, T b) {
 39         return b == 0 ? a : gcd(b, a % b); }
 40     template<class T> inline T fpw(T b, T i, T _m, T r = 1) {
 41         for(; i; i >>= 1, b = b * b % _m)
 42             if(i & 1) r = r * b % _m;
 43         return r; }
 44 };
 45 using namespace my_header;
 46 
 47 const int MAXN = 5e6 + 100;
 48 const int MAXV = 2e5 + 100;
 49 int n, col[MAXV], fa[MAXV], cnt[MAXV], ans[MAXV];
 50 int h[MAXV], nxt[MAXV], to[MAXV], e;
 51 
 52 #define lc ch[0]
 53 #define rc ch[1]
 54 struct SegNode {
 55     SegNode *ch[2];
 56     int s, c;
 57     void init() {
 58         s = c = 0;
 59         lc = rc = NULL;
 60     }
 61     void update();
 62     void merge(int, int, SegNode*);
 63     void insert(int, int, int);
 64 } node_pool[MAXN], *loc, *root[MAXV];
 65 SegNode* newSegNode() {
 66     loc->init();
 67     return loc++;
 68 }
 69 
 70 void SegNode::update() {
 71     s = (lc ? lc->s : 0) + (rc ? rc->s : 0);
 72 }
 73 void SegNode::insert(int l, int r, int v) {
 74     if (l == r) {
 75         ++c;
 76         s = (c != cnt[l]);
 77     } else {
 78         int m = (l + r) >> 1;
 79         if (v <= m) {
 80             if (!lc)
 81                 lc = newSegNode();
 82             lc->insert(l, m, v);
 83         } else {
 84             if (!rc)
 85                 rc = newSegNode();
 86             rc->insert(m + 1, r, v);
 87         }
 88         update();
 89     }
 90 }
 91 void SegNode::merge(int l, int r, SegNode* v) {
 92     if (l == r) {
 93         if (v) {
 94             c += v->c;
 95             s = (c != cnt[l]);
 96         }
 97     } else {
 98         int m = (l + r) >> 1;
 99         if (!lc)
100             lc = v->lc;
101         else if (v->lc)
102             lc->merge(l, m, v->lc);
103         if (!rc)
104             rc = v->rc;
105         else if (v->rc)
106             rc->merge(m + 1, r, v->rc);
107         update();
108     }
109 }
110 
111 void dfs(int u) {
112     int w = -1;
113     root[u] = newSegNode();
114     root[u]->insert(1, n, col[u]);
115     for (int i = h[u]; i != -1; i = nxt[i]) {
116         int v = to[i];
117         if (v != fa[u]) {
118             fa[v] = u;
119             dfs(v);
120             root[u]->merge(1, n, root[v]);
121         } else w = i >> 1;
122     }
123     if (w != -1)
124         ans[w] = root[u]->s;
125 }
126 
127 int main() {
128 #ifdef lol
129     freopen("1811.in", "r", stdin);
130     freopen("1811.out", "w", stdout);
131 #endif
132 
133     while (scanf("%d", &n) != EOF) {
134         for (int i = 1; i <= n; ++i)
135             cnt[i] = 0;
136         for (int i = 1; i <= n; ++i) {
137             fr(col[i]);
138             ++cnt[col[i]];
139         }
140         e = 0;
141         memset(h, -1, sizeof h);
142         for (int i = 1; i < n; ++i) {
143             int u, v;
144             fr(u, v);
145             to[e] = v, nxt[e] = h[u], h[u] = e++;
146             to[e] = u, nxt[e] = h[v], h[v] = e++;
147         }
148         loc = node_pool;
149         dfs(1);
150         for (int i = 0; i < n - 1; ++i)
151             wt(ans[i]);
152     }
153 
154     return 0;
155 }