fork(1) download
copy 
  1. #include <bits/stdc++.h>
  2.  
  3. using std::cerr;
  4.  
  5. struct top_tree_node {
  6. 	mutable top_tree_node* p = nullptr;
  7. 	top_tree_node* c[3] = {nullptr, nullptr, nullptr};
  8.  
  9. 	int d() const {
  10. 		assert(p);
  11. 		if (this == p->c[0]) {
  12. 			return 0;
  13. 		} else if (this == p->c[1]) {
  14. 			return 1;
  15. 		} else if (this == p->c[2]) {
  16. 			return 2;
  17. 		} else assert(false);
  18. 	}
  19. 	top_tree_node*& p_c() const { return p->c[d()]; } // p->c which points to you
  20.  
  21. 	// 3 types of verts: path edges, path verts, non-path edges
  22. 	bool is_path;
  23. 	bool is_vert;
  24.  
  25. 	bool r() const { return !p || p->is_path != is_path; }
  26.  
  27. 	int __id;
  28.  
  29. 	bool flip_path = false;
  30. 	int64_t subtree_size = 0;
  31. 	int64_t path_size = 0;
  32. 	int64_t subtree_lazy = 0;
  33. 	int64_t path_lazy = 0;
  34. 	int64_t tot_A = 0;
  35. 	int64_t own_A = 0;
  36. 	std::array<int64_t, 2> tot_sub_d = {0,0};
  37. 	std::array<int64_t, 2> tot_ans = {0,0};
  38.  
  39. 	void do_flip_path() {
  40. 		assert(is_path);
  41. 		flip_path ^= 1;
  42. 		std::swap(tot_sub_d[0], tot_sub_d[1]);
  43. 		std::swap(tot_ans[0], tot_ans[1]);
  44. 	}
  45.  
  46. 	void do_subtree_increment(int64_t v = 1) {
  47. 		//cerr << "doing subtree increment " << __id << ' ' << v << '\n';
  48. 		subtree_lazy += v;
  49. 		if (is_vert) own_A += v;
  50. 		tot_A += subtree_size * v;
  51. 		tot_ans[0] += tot_sub_d[0] * v;
  52. 		tot_ans[1] += tot_sub_d[1] * v;
  53. 	}
  54.  
  55. 	void do_path_increment(int64_t v = 1) {
  56. 		//cerr << "doing path increment " << __id << ' ' << v << '\n';
  57. 		assert(is_path);
  58. 		path_lazy += v;
  59. 		if (is_vert) own_A += v;
  60. 		tot_A += path_size * v;
  61. 		tot_ans[0] += (path_size * (path_size-1) / 2) * v;
  62. 		tot_ans[1] += (path_size * (path_size-1) / 2) * v;
  63. 	}
  64.  
  65. 	void downdate() {
  66. 		// fortunately, this doesn't interact with our ops
  67. 		if (flip_path) {
  68. 			assert(is_path);
  69. 			if (!is_vert) {
  70. 				if (c[0]) c[0]->do_flip_path();
  71. 				if (c[1]) c[1]->do_flip_path();
  72. 			}
  73. 			std::swap(c[0], c[1]);
  74. 			flip_path = false;
  75. 		}
  76.  
  77. 		if (subtree_lazy) {
  78. 			if (c[0]) c[0]->do_subtree_increment(subtree_lazy);
  79. 			if (c[1]) c[1]->do_subtree_increment(subtree_lazy);
  80. 			if (c[2]) c[2]->do_subtree_increment(subtree_lazy);
  81. 			subtree_lazy = 0;
  82. 		}
  83.  
  84. 		if (path_lazy) {
  85. 			assert(is_path);
  86. 			if (!is_vert) {
  87. 				//std::cerr << "downdating path lazy" << ' ' << __id << ' ' << path_lazy << '\n';
  88. 				assert(c[0] && c[1] && !c[2]);
  89. 				c[0]->do_path_increment(path_lazy);
  90. 				c[1]->do_path_increment(path_lazy);
  91. 			}
  92. 			path_lazy = 0;
  93. 		}
  94. 	}
  95.  
  96. 	void update() {
  97. 		assert(!flip_path && !subtree_lazy && !path_lazy);
  98. 		if (is_vert) assert(is_path);
  99.  
  100. 		subtree_size = is_vert;
  101. 		for (int z = 0; z <= 2; z++) {
  102. 			if (c[z]) subtree_size += c[z]->subtree_size;
  103. 		}
  104.  
  105. 		if (!is_path) {
  106. 			path_size = 0;
  107. 		} else {
  108. 			path_size = is_vert;
  109. 			if (!is_vert) {
  110. 				for (int z = 0; z < 2; z++) {
  111. 					if (c[z]) path_size += c[z]->path_size;
  112. 				}
  113. 			}
  114. 			assert(path_size >= 1);
  115. 		}
  116.  
  117. 		tot_A = is_vert ? own_A : 0;
  118. 		for (int z = 0; z <= 2; z++) {
  119. 			if (c[z]) tot_A += c[z]->tot_A;
  120. 		}
  121.  
  122. 		if (is_path && !is_vert) {
  123. 			assert(c[0] && c[1]);
  124. 			assert(!c[2]);
  125. 			for (int z = 0; z < 2; z++) {
  126. 				tot_sub_d[z] = c[z]->tot_sub_d[z] + c[z]->path_size * c[!z]->subtree_size + c[!z]->tot_sub_d[z];
  127. 				tot_ans[z] = c[z]->tot_ans[z] + c[z]->path_size * c[!z]->tot_A + c[!z]->tot_ans[z];
  128. 			}
  129. 		} else {
  130. 			tot_sub_d[0] = 0;
  131. 			tot_ans[0] = 0;
  132. 			for (int z = 0; z < 2; z++) {
  133. 				if (c[z]) {
  134. 					tot_sub_d[0] += c[z]->tot_sub_d[0];
  135. 					tot_ans[0] += c[z]->tot_ans[0];
  136. 				}
  137. 			}
  138.  
  139. 			if (c[2]) {
  140. 				assert(!is_path);
  141. 				tot_sub_d[0] += c[2]->subtree_size + c[2]->tot_sub_d[0];
  142. 				tot_ans[0] += c[2]->tot_A + c[2]->tot_ans[0];
  143. 			}
  144.  
  145. 			tot_sub_d[1] = tot_sub_d[0];
  146. 			tot_ans[1] = tot_ans[0];
  147. 		}
  148. 	}
  149.  
  150. 	void downdate_all() {
  151. 		if (p) p->downdate_all();
  152. 		downdate();
  153. 	}
  154.  
  155. 	void update_all() {
  156. 		update();
  157. 		if (p) p->update_all();
  158. 	}
  159.  
  160. private:
  161.  
  162. 	void rot() {
  163. 		assert(!is_vert);
  164. 		assert(!r());
  165. 		top_tree_node* pa = p;
  166. 		int x = d(); assert(x == 0 || x == 1);
  167. 		top_tree_node* ch = c[!x];
  168.  
  169. 		if (pa->p) pa->p_c() = this;
  170. 		this->p = pa->p;
  171.  
  172. 		pa->c[x] = ch;
  173. 		if (ch) ch->p = pa;
  174.  
  175. 		this->c[!x] = pa;
  176. 		pa->p = this;
  177.  
  178. 		pa->update();
  179. 	}
  180.  
  181. 	void rot_2(int c_d) {
  182. 		assert(!is_vert);
  183. 		assert(!r());
  184. 		assert(c[c_d]);
  185. 		assert(!c[c_d]->is_vert);
  186.  
  187. 		if (d() == c_d) {
  188. 			rot();
  189. 			return;
  190. 		}
  191.  
  192. 		top_tree_node* pa = p;
  193. 		int x = d(); assert(x == 0 || x == 1);
  194. 		assert(c_d == !x);
  195. 		top_tree_node* ch = c[c_d]->c[!x];
  196.  
  197. 		if (pa->p) pa->p_c() = this;
  198. 		this->p = pa->p;
  199.  
  200. 		pa->c[x] = ch;
  201. 		if (ch) ch->p = pa;
  202.  
  203. 		this->c[c_d]->c[!x] = pa;
  204. 		pa->p = this->c[c_d];
  205.  
  206. 		pa->update();
  207. 	}
  208.  
  209. 	void splay_dir(int x) {
  210. 		while (!r() && d() == x) {
  211. 			if (!p->r() && p->d() == x) {
  212. 				p->rot();
  213. 			}
  214. 			rot();
  215. 		}
  216. 	}
  217.  
  218. 	void splay_2(int c_d) {
  219. 		assert(!is_vert && is_path);
  220. 		assert(c[c_d] && !c[c_d]->is_vert);
  221. 		while (!r()) {
  222. 			if (!p->r()) {
  223. 				if (p->d() == d()) {
  224. 					p->rot();
  225. 				} else {
  226. 					rot_2(c_d);
  227. 				}
  228. 			}
  229. 			rot_2(c_d);
  230. 		}
  231. 	}
  232.  
  233. 	void splay_2() {
  234. 		assert(!is_vert && is_path);
  235. 		assert(!r());
  236. 		p->splay_2(d());
  237. 	}
  238.  
  239. 	void splay_vert() {
  240. 		assert(is_vert);
  241. 		if (r()) {
  242. 			return;
  243. 		}
  244. 		p->splay_dir(d());
  245. 		if (p->r()) {
  246. 			return;
  247. 		}
  248.  
  249. 		assert(p->d() != d());
  250. 		// we have a preference to be the left child
  251. 		if (d() == 1) {
  252. 			p->rot();
  253. 		}
  254. 		assert(d() == 0);
  255.  
  256. 		p->splay_2();
  257. 		assert(d() == 0);
  258. 		assert(p->d() == 1);
  259. 		assert(p->p->r());
  260. 	}
  261.  
  262. 	void splay() {
  263. 		assert(!is_vert);
  264. 		while (!r()) {
  265. 			if (!p->r()) {
  266. 				if (p->d() == d()) {
  267. 					p->rot();
  268. 				} else {
  269. 					rot();
  270. 				}
  271. 			}
  272. 			rot();
  273. 		}
  274. 	}
  275.  
  276. 	top_tree_node* cut_right() {
  277. 		assert(is_vert && is_path);
  278. 		splay_vert();
  279.  
  280. 		if (r() || d() == 1) {
  281. 			assert(r() || (d() == 1 && p->r()));
  282. 			assert(c[0] == nullptr);
  283. 			return nullptr;
  284. 		}
  285.  
  286. 		top_tree_node* pa = p;
  287. 		assert(pa->r() || (pa->d() == 1 && pa->p->r()));
  288. 		assert(!pa->is_vert);
  289. 		assert(pa->is_path);
  290. 		assert(pa->c[0] == this);
  291. 		assert(pa->c[2] == nullptr);
  292.  
  293. 		if (pa->p) pa->p_c() = this;
  294. 		this->p = pa->p;
  295.  
  296. 		pa->is_path = false;
  297. 		pa->c[2] = pa->c[1]; // don't need to change the parent
  298.  
  299. 		pa->c[0] = c[0]; if (c[0]) c[0]->p = pa;
  300. 		pa->c[1] = c[1]; if (c[1]) c[1]->p = pa;
  301.  
  302. 		c[0] = nullptr;
  303. 		c[1] = pa; pa->p = this;
  304. 		assert(c[2] == nullptr);
  305.  
  306. 		assert(c[0] == nullptr);
  307.  
  308. 		pa->update();
  309. 		return pa;
  310. 	}
  311.  
  312. 	top_tree_node* splice_non_path() {
  313. 		assert(!is_path);
  314. 		assert(!is_vert);
  315.  
  316. 		splay();
  317. 		assert(p && p->is_vert && p->is_path);
  318. 		p->cut_right();
  319.  
  320. 		if (!p->is_path) rot();
  321. 		assert(p && p->is_vert && p->is_path);
  322. 		assert(p->r() || (p->d() == 1 && p->p->r()));
  323. 		assert(p->c[d()] == this && p->c[!d()] == nullptr);
  324.  
  325. 		top_tree_node* pa = p;
  326.  
  327. 		if (pa->p) pa->p_c() = this;
  328. 		this->p = pa->p;
  329.  
  330. 		pa->c[0] = c[0]; if (c[0]) c[0]->p = pa;
  331. 		pa->c[1] = c[1]; if (c[1]) c[1]->p = pa;
  332.  
  333. 		assert(c[2] && c[2]->is_path);
  334. 		c[1] = c[2]; // don't need to change parent
  335. 		c[0] = pa; pa->p = this;
  336. 		c[2] = nullptr;
  337.  
  338. 		is_path = true;
  339.  
  340. 		pa->update();
  341. 		return pa;
  342. 	}
  343.  
  344. 	void splice_all(top_tree_node*& res) {
  345. 		if (!is_path) {
  346. 			res = splice_non_path();
  347. 		}
  348. 		assert(is_path);
  349. 		if (!p) return;
  350. 		p->splice_all(res);
  351. 	}
  352.  
  353. public:
  354. 	top_tree_node* expose() {
  355. 		assert(is_vert);
  356. 		downdate_all();
  357.  
  358. 		top_tree_node* res = nullptr;
  359. 		splice_all(res);
  360.  
  361. 		cut_right();
  362.  
  363. 		update_all();
  364.  
  365. 		return res;
  366. 	}
  367.  
  368. 	void meld_path_end() {
  369. 		assert(!p);
  370. 		top_tree_node* rt = this;
  371. 		while (true) {
  372. 			rt->downdate();
  373. 			if (rt->is_vert) break;
  374. 			rt = rt->c[1];
  375. 		}
  376. 		assert(rt->is_vert);
  377. 		rt->splay_vert();
  378. 		if (rt->c[0] && rt->c[1]) {
  379. 			top_tree_node* ch = rt->c[1];
  380. 			while (true) {
  381. 				ch->downdate();
  382. 				if (!ch->c[0]) break;
  383. 				ch = ch->c[0];
  384. 			}
  385. 			ch->splay();
  386. 			assert(ch->c[0] == nullptr);
  387.  
  388. 			ch->c[0] = rt->c[0];
  389. 			ch->c[0]->p = ch;
  390.  
  391. 			rt->c[0] = nullptr;
  392.  
  393. 			ch->update();
  394. 		} else if (rt->c[0]) {
  395. 			rt->c[1] = rt->c[0];
  396. 			rt->c[0] = nullptr;
  397. 		}
  398. 		assert(rt->c[0] == nullptr);
  399. 		rt->update_all();
  400. 	}
  401.  
  402. 	void make_root() {
  403. 		expose();
  404.  
  405. 		top_tree_node* rt = this;
  406. 		while (rt->p) {
  407. 			assert(rt->d() == 1);
  408. 			rt = rt->p;
  409. 		}
  410. 		rt->do_flip_path();
  411. 		rt->meld_path_end();
  412.  
  413. 		expose();
  414.  
  415. 		assert(!p);
  416. 	}
  417.  
  418. 	// link v2 as a child of v1 with edge e
  419. 	friend void link(top_tree_node* e, top_tree_node* v1, top_tree_node* v2) {
  420. 		assert(e && v1 && v2);
  421. 		assert(!e->c[0] && !e->c[1] && !e->c[2]);
  422. 		v1->expose(); while (v1->p) v1 = v1->p;
  423. 		v2->make_root();
  424.  
  425. 		assert(!v1->p);
  426. 		assert(!v2->p);
  427.  
  428. 		e->is_path = true, e->is_vert = false;
  429. 		e->c[0] = v1;
  430. 		v1->p = e;
  431. 		e->c[1] = v2;
  432. 		v2->p = e;
  433. 		e->update();
  434. 	}
  435.  
  436. 	friend std::pair<top_tree_node*, top_tree_node*> cut(top_tree_node* e) {
  437. 		assert(!e->p);
  438. 		assert(e->is_path);
  439. 		assert(!e->is_vert);
  440.  
  441. 		e->downdate();
  442.  
  443. 		top_tree_node* l = e->c[0];
  444. 		top_tree_node* r = e->c[1];
  445. 		assert(l && r);
  446.  
  447. 		e->c[0] = e->c[1] = nullptr;
  448. 		l->p = r->p = nullptr;
  449.  
  450. 		assert(e->c[2] == nullptr);
  451.  
  452. 		l->meld_path_end();
  453.  
  454. 		return {l, r};
  455. 	}
  456.  
  457. 	friend top_tree_node* get_path(top_tree_node* a, top_tree_node* b) {
  458. 		assert(a->is_vert && b->is_vert);
  459. 		a->make_root();
  460. 		b->expose();
  461. 		if (a == b) {
  462. 			assert(!b->p);
  463. 			return b;
  464. 		}
  465. 		assert(!b->p->p);
  466. 		return b->p;
  467. 	}
  468.  
  469. 	friend top_tree_node* get_subtree(top_tree_node* rt, top_tree_node* n) {
  470. 		rt->make_root();
  471. 		n->expose();
  472. 		return n;
  473. 	}
  474. };
  475.  
  476. int main() {
  477. 	using namespace std;
  478. 	ios_base::sync_with_stdio(false), cin.tie(nullptr);
  479.  
  480. 	int N; cin >> N;
  481. 	vector<top_tree_node> nodes(2*N-1);
  482. 	for (int i = 0; i < 2 * N - 1; i++) {
  483. 		nodes[i].__id = i;
  484. 	}
  485.  
  486. 	for (int i = 0; i < N; i++) {
  487. 		top_tree_node* n = &nodes[i];
  488. 		n->is_path = n->is_vert = true;
  489. 		n->update();
  490. 	}
  491.  
  492. 	for (int e = 0; e < N-1; e++) {
  493. 		int u, v; cin >> u >> v; u--, v--;
  494. 		top_tree_node* a = &nodes[u];
  495. 		top_tree_node* b = &nodes[v];
  496. 		link(&nodes[N+e], a, b);
  497. 	}
  498.  
  499. 	int Q; cin >> Q;
  500. 	while (Q--) {
  501. 		int op; cin >> op;
  502. 		if (op == 1) {
  503. 			int u, v; cin >> u >> v; u--, v--;
  504. 			auto sub = get_subtree(&nodes[u], &nodes[v]);
  505. 			sub->do_subtree_increment();
  506. 			sub->downdate();
  507. 			sub->update_all();
  508. 		} else if (op == 2) {
  509. 			int u, v; cin >> u >> v; u--, v--;
  510. 			auto pth = get_path(&nodes[u], &nodes[v]);
  511. 			pth->do_path_increment();
  512. 			pth->downdate();
  513. 			pth->update_all();
  514. 		} else if (op == 3) {
  515. 			int v; cin >> v; v--;
  516. 			nodes[v].make_root();
  517. 			cout << nodes[v].tot_ans[0] << '\n';
  518. 		} else assert(false);
  519.  
  520. 		/*
  521. 		cerr << "dumping tree" << '\n';
  522. 		for (int z = 0; z < 2*N-1; z++) {
  523. 			cerr << "node " << z << '\n';
  524. 			cerr << "par: " << (nodes[z].p ? nodes[z].p->__id : -1) << '\n';
  525. 			cerr << "subtree_size: " << nodes[z].subtree_size << '\n';
  526. 			cerr << "path_size: " << nodes[z].path_size << '\n';
  527. 			cerr << "subtree_lazy: " << nodes[z].subtree_lazy << '\n';
  528. 			cerr << "path_lazy: " << nodes[z].path_lazy << '\n';
  529. 			cerr << "tot_A: " << nodes[z].tot_A << '\n';
  530. 			cerr << "own_A: " << nodes[z].own_A << '\n';
  531. 			cerr << "tot_sub_d[0]: " << nodes[z].tot_sub_d[0] << '\n';
  532. 			cerr << "tot_sub_d[1]: " << nodes[z].tot_sub_d[1] << '\n';
  533. 			cerr << "tot_ans[0]: " << nodes[z].tot_ans[0] << '\n';
  534. 			cerr << "tot_ans[1]: " << nodes[z].tot_ans[1] << '\n';
  535. 		}
  536. 		cerr << '\n';
  537. 		*/
  538. 	}
  539.  
  540. 	return 0;
  541. }
  542.  
Success #stdin #stdout 0s 4768KB
comments (?)
stdin
copy 
5
4 2
2 5
1 5
1 3
5
2 2 4
3 4
2 1 5
2 5 5
3 2
stdout
copy 
1
5
https://ideone.com/B03DBB
language:
C++14 (gcc 8.3)
created:
3 years ago
visibility:
public

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!