fork download
copy 
  1. #include <bits/stdc++.h>
  2. using namespace std;
  3.  
  4. using ll = long long;
  5. using ld = long double;
  6. using db = double; 
  7. using str = string; // yay python!
  8.  
  9. using pi = pair<int,int>;
  10. using pl = pair<ll,ll>;
  11. using pd = pair<db,db>;
  12.  
  13. using vi = vector<int>;
  14. using vb = vector<bool>;
  15. using vl = vector<ll>;
  16. using vd = vector<db>; 
  17. using vs = vector<str>;
  18. using vpi = vector<pi>;
  19. using vpl = vector<pl>; 
  20. using vpd = vector<pd>;
  21.  
  22. #define tcT template<class T
  23. // ^ lol this makes everything look weird but I'll try it
  24. tcT> using V = vector<T>; 
  25. tcT, size_t SZ> using AR = array<T,SZ>; 
  26.  
  27. // pairs
  28. #define mp make_pair
  29. #define f first
  30. #define s second
  31.  
  32. // vectors
  33. #define sz(x) (int)(x).size()
  34. #define all(x) begin(x), end(x)
  35. #define rall(x) (x).rbegin(), (x).rend() 
  36. #define sor(x) sort(all(x)) 
  37. #define rsz resize
  38. #define ins insert 
  39. #define ft front() 
  40. #define bk back()
  41. #define pf push_front 
  42. #define pb push_back
  43. #define eb emplace_back 
  44. #define lb lower_bound 
  45. #define ub upper_bound 
  46.  
  47. // loops
  48. #define FOR(i,a,b) for (int i = (a); i < (b); ++i)
  49. #define F0R(i,a) FOR(i,0,a)
  50. #define ROF(i,a,b) for (int i = (b)-1; i >= (a); --i)
  51. #define R0F(i,a) ROF(i,0,a)
  52. #define trav(a,x) for (auto& a: x)
  53.  
  54. const int MOD = 1e9+7; // 998244353;
  55. const int MX = 50005;
  56. const ll INF = 1e18; // not too close to LLONG_MAX
  57. const ld PI = acos((ld)-1);
  58. const int xd[4] = {1,0,-1,0}, yd[4] = {0,1,0,-1}; // for every grid problem!!
  59. mt19937 rng((uint32_t)chrono::steady_clock::now().time_since_epoch().count()); 
  60.  
  61. // helper funcs
  62. constexpr int pct(int x) { return __builtin_popcount(x); } // # of bits set
  63. constexpr int bits(int x) { return 31-__builtin_clz(x); } // floor(log2(x)) 
  64. ll cdiv(ll a, ll b) { return a/b+((a^b)>0&&a%b); } // divide a by b rounded up
  65. ll fdiv(ll a, ll b) { return a/b-((a^b)<0&&a%b); } // divide a by b rounded down
  66. ll half(ll x) { return fdiv(x,2); }
  67.  
  68. tcT> bool ckmin(T& a, const T& b) {
  69. 	return b < a ? a = b, 1 : 0; } // set a = min(a,b)
  70. tcT> bool ckmax(T& a, const T& b) {
  71. 	return a < b ? a = b, 1 : 0; }
  72.  
  73. #define tcTU tcT, class U
  74. tcTU> T fstTrue(T lo, T hi, U f) { 
  75. 	// note: if (lo+hi)/2 is used instead of half(lo+hi) 
  76. 	// then this will loop infinitely when lo=hi
  77. 	hi ++; assert(lo <= hi); // assuming f is increasing
  78. 	while (lo < hi) { // find first index such that f is true 
  79. 		T mid = half(lo+hi);
  80. 		f(mid) ? hi = mid : lo = mid+1; 
  81. 	} 
  82. 	return lo;
  83. }
  84. tcTU> T lstTrue(T lo, T hi, U f) {
  85. 	lo --; assert(lo <= hi); // assuming f is decreasing
  86. 	while (lo < hi) { // find first index such that f is true 
  87. 		T mid = half(lo+hi+1);
  88. 		f(mid) ? lo = mid : hi = mid-1;
  89. 	} 
  90. 	return lo;
  91. }
  92. tcT> void remDup(vector<T>& v) { // sort and remove duplicates
  93. 	sort(all(v)); v.erase(unique(all(v)),end(v)); }
  94. tcTU> void erase(T& t, const U& u) { // don't erase
  95. 	auto it = t.find(u); assert(it != end(t));
  96. 	t.erase(u); } // element that doesn't exist from (multi)set
  97.  
  98. // INPUT
  99. #define tcTUU tcT, class ...U
  100. tcT> void re(complex<T>& c);
  101. tcTU> void re(pair<T,U>& p);
  102. tcT> void re(vector<T>& v);
  103. tcT, size_t SZ> void re(AR<T,SZ>& a);
  104.  
  105. tcT> void re(T& x) { cin >> x; }
  106. void re(db& d) { str t; re(t); d = stod(t); }
  107. void re(ld& d) { str t; re(t); d = stold(t); }
  108. tcTUU> void re(T& t, U&... u) { re(t); re(u...); }
  109.  
  110. tcT> void re(complex<T>& c) { T a,b; re(a,b); c = {a,b}; }
  111. tcTU> void re(pair<T,U>& p) { re(p.f,p.s); }
  112. tcT> void re(vector<T>& x) { trav(a,x) re(a); }
  113. tcT, size_t SZ> void re(AR<T,SZ>& x) { trav(a,x) re(a); }
  114.  
  115. // TO_STRING
  116. #define ts to_string
  117. str ts(char c) { return str(1,c); }
  118. str ts(const char* s) { return (str)s; }
  119. str ts(str s) { return s; }
  120. str ts(bool b) { 
  121. 	#ifdef LOCAL
  122. 		return b ? "true" : "false"; 
  123. 	#else 
  124. 		return ts((int)b);
  125. 	#endif
  126. }
  127. tcT> str ts(complex<T> c) { 
  128. 	stringstream ss; ss << c; return ss.str(); }
  129. str ts(vector<bool> v) {
  130. 	str res = "{"; F0R(i,sz(v)) res += char('0'+v[i]);
  131. 	res += "}"; return res; }
  132. template<size_t SZ> str ts(bitset<SZ> b) {
  133. 	str res = ""; F0R(i,SZ) res += char('0'+b[i]);
  134. 	return res; }
  135. tcTU> str ts(pair<T,U> p);
  136. tcT> str ts(T v) { // containers with begin(), end()
  137. 	#ifdef LOCAL
  138. 		bool fst = 1; str res = "{";
  139. 		for (const auto& x: v) {
  140. 			if (!fst) res += ", ";
  141. 			fst = 0; res += ts(x);
  142. 		}
  143. 		res += "}"; return res;
  144. 	#else
  145. 		bool fst = 1; str res = "";
  146. 		for (const auto& x: v) {
  147. 			if (!fst) res += " ";
  148. 			fst = 0; res += ts(x);
  149. 		}
  150. 		return res;
  151.  
  152. 	#endif
  153. }
  154. tcTU> str ts(pair<T,U> p) {
  155. 	#ifdef LOCAL
  156. 		return "("+ts(p.f)+", "+ts(p.s)+")"; 
  157. 	#else
  158. 		return ts(p.f)+" "+ts(p.s);
  159. 	#endif
  160. }
  161.  
  162. // OUTPUT
  163. tcT> void pr(T x) { cout << ts(x); }
  164. tcTUU> void pr(const T& t, const U&... u) { 
  165. 	pr(t); pr(u...); }
  166. void ps() { pr("\n"); } // print w/ spaces
  167. tcTUU> void ps(const T& t, const U&... u) { 
  168. 	pr(t); if (sizeof...(u)) pr(" "); ps(u...); }
  169.  
  170. // DEBUG
  171. void DBG() { cerr << "]" << endl; }
  172. tcTUU> void DBG(const T& t, const U&... u) {
  173. 	cerr << ts(t); if (sizeof...(u)) cerr << ", ";
  174. 	DBG(u...); }
  175. #ifdef LOCAL // compile with -DLOCAL, chk -> fake assert
  176. 	#define dbg(...) cerr << "Line(" << __LINE__ << ") -> [" << #__VA_ARGS__ << "]: [", DBG(__VA_ARGS__)
  177. 	#define chk(...) if (!(__VA_ARGS__)) cerr << "Line(" << __LINE__ << ") -> function(" \
  178. 		 << __FUNCTION__  << ") -> CHK FAILED: (" << #__VA_ARGS__ << ")" << "\n", exit(0);
  179. #else
  180. 	#define dbg(...) 0
  181. 	#define chk(...) 0
  182. #endif
  183.  
  184. // FILE I/O
  185. void setIn(str s) { freopen(s.c_str(),"r",stdin); }
  186. void setOut(str s) { freopen(s.c_str(),"w",stdout); }
  187. void unsyncIO() { cin.tie(0)->sync_with_stdio(0); }
  188. void setIO(str s = "") {
  189. 	unsyncIO();
  190. 	// cin.exceptions(cin.failbit); 
  191. 	// throws exception when do smth illegal
  192. 	// ex. try to read letter into int
  193. 	if (sz(s)) { setIn(s+".in"), setOut(s+".out"); } // for USACO
  194. }
  195.  
  196. /**
  197.  * Description: 1D range increment and sum query.
  198.  * Source: USACO Counting Haybales
  199.  * Verification: SPOJ Horrible
  200.  */
  201.  
  202. template<class T, int SZ> struct LazySeg { 
  203. 	static_assert(pct(SZ) == 1); // SZ must be power of 2
  204. 	const T ID = {{0,0}}; 
  205. 	T comb(T a, T b) { return {a[0]+b[0],a[1]+b[1]}; }
  206. 	T seg[2*SZ];
  207. 	int lazy[2*SZ]; 
  208. 	LazySeg() { 
  209. 		F0R(i,2*SZ) seg[i] = ID; 
  210. 	}
  211. 	void push(int ind, int L, int R) { /// modify values for current node
  212. 		if (!lazy[ind]) return;
  213. 		swap(seg[ind][0],seg[ind][1]);
  214. 		// seg[ind] += (R-L+1)*lazy[ind]; // dependent on operation
  215. 		if (L != R) F0R(i,2) lazy[2*ind+i] ^= lazy[ind]; /// prop to children
  216. 		lazy[ind] = 0; 
  217. 	} // recalc values for current node
  218. 	void pull(int ind) { seg[ind] = comb(seg[2*ind],seg[2*ind+1]); }
  219. 	void build() { ROF(i,1,SZ) pull(i); }
  220. 	void upd(int lo,int hi,int ind=1,int L=0, int R=SZ-1) {
  221. 		push(ind,L,R); if (hi < L || R < lo) return;
  222. 		if (lo <= L && R <= hi) { 
  223. 			lazy[ind] = 1; push(ind,L,R); return; }
  224. 		int M = (L+R)/2; 
  225. 		upd(lo,hi,2*ind,L,M); 
  226. 		upd(lo,hi,2*ind+1,M+1,R); 
  227. 		pull(ind);
  228. 	}
  229. 	T query(int lo, int hi, int ind=1, int L=0, int R=SZ-1) {
  230. 		push(ind,L,R); if (lo > R || L > hi) return ID;
  231. 		if (lo <= L && R <= hi) return seg[ind];
  232. 		int M = (L+R)/2; 
  233. 		return comb(query(lo,hi,2*ind,L,M),query(lo,hi,2*ind+1,M+1,R));
  234. 	}
  235. };
  236.  
  237. LazySeg<AR<int,2>,1<<16> L;
  238.  
  239. /**
  240.  * Description: Heavy-Light Decomposition, add val to verts 
  241.  	* and query sum in path/subtree.
  242.  * Time: any tree path is split into $O(\log N)$ parts
  243.  * Source: http://c...content-available-to-author-only...s.com/blog/entry/22072, https://c...content-available-to-author-only...s.com/blog/entry/53170
  244.  * Verification: *
  245.  */
  246.  
  247. // #include "../../data-structures/1D Range Queries (9.2)/LazySeg (15.2).h"
  248.  
  249. const int SZ = 50005;
  250.  
  251. vi adj[SZ];
  252. int par[SZ], root[SZ], depth[SZ], sz[SZ], ti;
  253. int pos[SZ]; vi rpos; // rpos not used, but could be useful
  254. ll stor[MX];
  255.  
  256. template<int SZ, bool VALS_IN_EDGES> struct HLD { 
  257. 	void ae(int x, int y) { adj[x].pb(y), adj[y].pb(x); }
  258. 	void dfsSz(int x) { 
  259. 		sz[x] = 1; 
  260. 		trav(y,adj[x]) {
  261. 			par[y] = x; depth[y] = depth[x]+1;
  262. 			adj[y].erase(find(all(adj[y]),x)); // remove parent from adj list
  263. 			dfsSz(y); sz[x] += sz[y];
  264. 			if (sz[y] > sz[adj[x][0]]) swap(y,adj[x][0]);
  265. 		}
  266. 	}
  267. 	void dfsHld(int x) {
  268. 		pos[x] = ti++; rpos.pb(x);
  269. 		trav(y,adj[x]) {
  270. 			root[y] = (y == adj[x][0] ? root[x] : y);
  271. 			dfsHld(y); }
  272. 	}
  273. 	void init() {
  274. 		int R = 0;
  275. 		par[R] = depth[R] = ti = 0; dfsSz(R); 
  276. 		root[R] = R; dfsHld(R); 
  277. 	}
  278. 	int lca(int x, int y) {
  279. 		for (; root[x] != root[y]; y = par[root[y]])
  280. 			if (depth[root[x]] > depth[root[y]]) swap(x,y);
  281. 		return depth[x] < depth[y] ? x : y;
  282. 	}
  283. 	/// int dist(int x, int y) { // # edges on path
  284. 	/// 	return depth[x]+depth[y]-2*depth[lca(x,y)]; }
  285. 	// LazySeg<ll,SZ> tree; // segtree for sum
  286. 	// template <class BinaryOp>
  287. 	// void processPath(int x, int y, BinaryOp op) {
  288. 	// 	for (; root[x] != root[y]; y = par[root[y]]) {
  289. 	// 		if (depth[root[x]] > depth[root[y]]) swap(x,y);
  290. 	// 		op(pos[root[y]],pos[y]); }
  291. 	// 	if (depth[x] > depth[y]) swap(x,y);
  292. 	// 	op(pos[x]+VALS_IN_EDGES,pos[y]); 
  293. 	// }
  294. 	// void modifyPath(int x, int y, int v) { 
  295. 	// 	processPath(x,y,[this,&v](int l, int r) { 
  296. 	// 		tree.upd(l,r,v); }); }
  297. 	// ll queryPath(int x, int y) { 
  298. 	// 	ll res = 0; processPath(x,y,[this,&res](int l, int r) { 
  299. 	// 		res += tree.query(l,r); });
  300. 	// 	return res; }
  301. 	// void modifySubtree(int x, int v) { 
  302. 	// 	tree.upd(pos[x]+VALS_IN_EDGES,pos[x]+sz[x]-1,v); }
  303. };
  304.  
  305. ll sq(int x) {
  306. 	ll num = L.query(pos[x],pos[x]+sz[x]-1)[1];
  307. 	return num*(num-1)/2;
  308. }
  309.  
  310. HLD<50005,0> H;
  311. int N,Q,A[MX];
  312. // vi child[MX];
  313.  
  314. void flip(int x) {
  315. 	for (int y = x; 0 != root[y]; y = par[root[y]]) {
  316. 		int p = par[root[y]];
  317. 		stor[p] -= sq(root[y]);
  318. 	}
  319. 	for (int y = x; ; y = par[root[y]]) {
  320. 		if (root[y] == 0) {
  321. 			L.upd(pos[0],pos[y]); 
  322. 			break;
  323. 		} else {
  324. 			L.upd(pos[root[y]],pos[y]); 
  325. 		}
  326. 	}
  327. 	for (int y = x; 0 != root[y]; y = par[root[y]]) {
  328. 		int p = par[root[y]];
  329. 		stor[p] += sq(root[y]);
  330. 	}
  331. }
  332.  
  333. int main() {
  334. 	setIO(); re(N,Q);
  335. 	F0R(i,N) re(A[i]);
  336. 	F0R(i,N-1) {
  337. 		int u,v; re(u,v); u--,v--;
  338. 		H.ae(u,v);
  339. 	}
  340. 	H.init();
  341. 	F0R(i,N) {
  342. 		L.seg[(1<<16)+pos[i]][A[i]] = 1;
  343. 	} 
  344. 	L.build();
  345. 	F0R(i,N) {
  346. 		FOR(j,1,sz(adj[i])) stor[i] += sq(adj[i][j]);
  347. 	}
  348. 	// F0R(i,N) dbg(i,sq(i));
  349. 	F0R(i,Q) {
  350. 		int o,x; re(o,x); x --;
  351. 		if (o == 1) {
  352. 			flip(x);
  353. 		} else {
  354. 			ll ans = sq(x)-stor[x];
  355. 			if (sz(adj[x])) ans -= sq(adj[x][0]);
  356. 			ps(ans);
  357. 		}
  358. 	}
  359. 	// HLD? block tree?
  360.  
  361. 	// you should actually read the stuff at the bottom
  362. }
  363.  
  364. /* stuff you should look for
  365. 	* int overflow, array bounds
  366. 	* special cases (n=1?)
  367. 	* do smth instead of nothing and stay organized
  368. 	* WRITE STUFF DOWN
  369. 	* DON'T GET STUCK ON ONE APPROACH
  370. */
Success #stdin #stdout 0s 5748KB
comments (?)
stdin
copy 
Standard input is empty
stdout
copy 
Standard output is empty
https://ideone.com/xXunf6
TROC 14 F
language:
C++14 (gcc 8.3)
created:
3 years ago
visibility:
secret

Share or Embed source code

Discover > Sphere Engine API

The brand new service which powers Ideone!

Discover > IDE Widget

Widget for compiling and running the source code in a web browser!