fork(1) download
copy 
  1. #include <stdio.h>
  2. #include <stdlib.h>
  3.  
  4. #define RED    	1
  5. #define BLACK	2
  6.  
  7. struct node {
  8. 	long long key;
  9. 	struct node *left, *right, *p;
  10. 	long long color;
  11. 	long long size;
  12. };
  13.  
  14. typedef struct node *NODEPTR;
  15. struct node NIL;
  16. NODEPTR NILPTR = &NIL;
  17.  
  18. void inorder(NODEPTR x) {
  19. 	if (x != NILPTR) {
  20. 		inorder(x->left);
  21. 		printf("(%lld,%lld) ", x->key, x->size);
  22. 		inorder(x->right);
  23. 	}
  24. 	//putchar('\n');
  25. }
  26.  
  27. void preorder(NODEPTR x) {
  28. 	if (x != NILPTR) {
  29. 		printf("%lld ", x->key);
  30. 		preorder(x->left);
  31. 		preorder(x->right);
  32. 	}
  33. 	//putchar('\n');
  34. }
  35.  
  36. NODEPTR search(NODEPTR root, long long k) {
  37. 	if (root == NILPTR || root->key == k)
  38. 		return root;
  39. 	if (k < root->key)
  40. 		return search(root->left, k);
  41. 	else
  42. 		return search(root->right, k);
  43. }
  44.  
  45. NODEPTR minimum(NODEPTR root) {
  46. 	while (root->left != NILPTR)
  47. 		root = root->left;
  48. 	return root;
  49. }
  50.  
  51. NODEPTR maximum(NODEPTR root) {
  52. 	while (root->right != NILPTR)
  53. 		root = root->right;
  54. 	return root;
  55. }
  56.  
  57. NODEPTR successor(NODEPTR root, long long x) {
  58. 	NODEPTR temp = search(root, x);
  59. 	if (temp == NILPTR) {
  60. 		printf("%lld not in tree\n", x);
  61. 		return temp;
  62. 	}
  63. 	if (temp->right != NILPTR)
  64. 		return minimum(temp->right);
  65. 	NODEPTR y = temp->p;
  66. 	while (y != NILPTR && temp == y->right) {
  67. 		temp = y;
  68. 		y = y->p;
  69. 	}
  70. 	return y;
  71. }
  72.  
  73. NODEPTR predecessor(NODEPTR root, long long x) {
  74. 	NODEPTR temp = search(root, x);
  75. 	if (temp == NILPTR) {
  76. 		printf("%lld not in tree\n", x);
  77. 		return temp;
  78. 	}
  79. 	if (temp->left != NILPTR)
  80. 		return maximum(temp->left);
  81. 	NODEPTR y = temp->p;
  82. 	while (y != NILPTR && temp == y->left) {
  83. 		temp = y;
  84. 		y = y->p;
  85. 	}
  86. 	return y;
  87. }
  88. void leftrotate(NODEPTR *treeroot, NODEPTR x) {
  89. 	NODEPTR y = x->right;
  90. 	x->right = y->left;
  91. 	if (y->left != NILPTR)
  92. 		y->left->p = x;
  93. 	y->p = x->p;
  94. 	if (x->p == NILPTR)
  95. 		*treeroot = y;
  96. 	else if (x->p->left == x)
  97. 		x->p->left = y;
  98. 	else
  99. 		x->p->right = y;
  100. 	y->left = x;
  101. 	x->p = y;
  102. 	y->size = x->size;
  103. 	x->size = x->left->size + x->right->size +1;
  104. }
  105.  
  106. void rightrotate(NODEPTR *treeroot, NODEPTR y) {
  107. 	NODEPTR x = y->left;
  108. 	y->left = x->right;
  109. 	if (x->right != NILPTR)
  110. 		x->right->p = y;
  111. 	x->p = y->p;
  112. 	if (y->p == NILPTR)
  113. 		*treeroot = x;
  114. 	else if (y->p->left == y)
  115. 		y->p->left = x;
  116. 	else
  117. 		y->p->right = x;
  118. 	x->right = y;
  119. 	y->p = x;
  120. 	x->size = y->size;
  121. 	y->size = y->left->size + y->right->size + 1;
  122. }
  123.  
  124. void rbinsertfixup(NODEPTR *treeroot, NODEPTR z) {
  125. 	while (z->p->color == RED) {
  126. 		if (z->p == z->p->p->left) {
  127. 			NODEPTR y = z->p->p->right;
  128. 			if (y->color == RED) {
  129. 				z->p->color = BLACK;
  130. 				y->color = BLACK;
  131. 				z->p->p->color = RED;
  132. 				z = z->p->p;
  133. 			}
  134. 			else {
  135. 				if (z == z->p->right) {
  136. 					z = z->p;
  137. 					leftrotate(treeroot,z);
  138. 				}
  139. 				z->p->color = BLACK;
  140. 				z->p->p->color = RED;
  141. 				rightrotate(treeroot,z->p->p);
  142. 			}
  143. 		}
  144. 		else {
  145. 			NODEPTR y = z->p->p->left;
  146. 			if (y->color == RED) {
  147. 				z->p->color = BLACK;
  148. 				y->color = BLACK;
  149. 				z->p->p->color = RED;
  150. 				z = z->p->p;
  151. 			}
  152. 			else {
  153. 				if (z == z->p->left) {
  154. 					z = z->p;
  155. 					rightrotate(treeroot,z);
  156. 				}
  157. 				z->p->color = BLACK;
  158. 				z->p->p->color = RED;
  159. 				leftrotate(treeroot,z->p->p);
  160. 			}
  161. 		}
  162. 	}
  163. 	(*treeroot)->color = BLACK;
  164. }
  165.  
  166. void rbinsert(NODEPTR *treeroot, long long z) {
  167. 	NODEPTR Z = (NODEPTR) malloc(sizeof(struct node));
  168. 	Z->key = z;
  169. 	NODEPTR y = NILPTR;
  170. 	NODEPTR x = *treeroot;
  171. 	while (x != NILPTR) {
  172. 		y = x;
  173. 		y->size++;
  174. 		if (Z->key < x->key)
  175. 			x = x->left;
  176. 		else if (Z->key > x->key)
  177. 			x = x->right;
  178. 		else {
  179. 			NODEPTR temp = x;
  180. 			while (x != NILPTR) {
  181. 				x->size--;
  182. 				x = x->p;
  183. 			}
  184. 			free(Z);
  185. 			return;
  186. 		}
  187. 	}
  188. 	Z->p = y;
  189. 	if (y == NILPTR)
  190. 		*treeroot = Z;
  191. 	else if (Z->key < y->key)
  192. 		y->left = Z;
  193. 	else
  194. 		y->right = Z;
  195. 	Z->left = NILPTR;
  196. 	Z->right = NILPTR;
  197. 	Z->color = RED;
  198. 	Z->size = 1;
  199. 	rbinsertfixup(treeroot,Z);
  200. }
  201.  
  202. void rbtransplant(NODEPTR *treeroot, NODEPTR u, NODEPTR v) {
  203. 	if (u->p == NILPTR)
  204. 		*treeroot = v;
  205. 	else if (u == u->p->left)
  206. 		u->p->left = v;
  207. 	else
  208. 		u->p->right = v;
  209. 	v->p = u->p;
  210. }
  211.  
  212. void rbdeletefixup(NODEPTR *treeroot, NODEPTR x) {
  213. 	while (x != *treeroot && x->color == BLACK) {
  214. 		if (x == x->p->left) {
  215. 			NODEPTR w = x->p->right;
  216. 			if (w->color == RED) {
  217. 				w->color = BLACK;
  218. 				x->p->color = RED;
  219. 				leftrotate(treeroot,x->p);
  220. 				w = x->p->right;
  221. 			}
  222. 			if (w->left->color == BLACK && w->right->color == BLACK) {
  223. 				w->color = RED;
  224. 				x = x->p;
  225. 			}
  226. 			else {
  227. 			 	if (w->right->color == BLACK) {
  228. 					w->left->color = BLACK;
  229. 					w->color = RED;
  230. 					rightrotate(treeroot,w);
  231. 					w = x->p->right;
  232. 				}
  233. 				w->color = x->p->color;
  234. 				x->p->color = BLACK;
  235. 				w->right->color = BLACK;
  236. 				leftrotate(treeroot,x->p);
  237. 				x = *treeroot;
  238. 			}
  239. 		}
  240. 		else {
  241. 			NODEPTR w = x->p->left;
  242. 			if (w->color == RED) {
  243. 				w->color = BLACK;
  244. 				x->p->color = RED;
  245. 				rightrotate(treeroot,x->p);
  246. 				w = x->p->left;
  247. 			}
  248. 			if (w->left->color == BLACK && w->right->color == BLACK) {
  249. 				w->color = RED;
  250. 				x = x->p;
  251. 			}
  252. 			else {
  253. 				if (w->left->color == BLACK) {
  254. 					w->right->color = BLACK;
  255. 					w->color = RED;
  256. 					leftrotate(treeroot,w);
  257. 					w = x->p->left;
  258. 				}
  259. 				w->color = x->p->color;
  260. 				x->p->color = BLACK;
  261. 				w->left->color = BLACK;
  262. 				rightrotate(treeroot,x->p);
  263. 				x = *treeroot;
  264. 			}
  265. 		}
  266. 	}
  267. 	x->color = BLACK;
  268. }
  269.  
  270. void rbdelete(NODEPTR *treeroot, long long z) {
  271. 	NODEPTR Z = search(*treeroot, z);
  272. 	if (Z == NILPTR) {
  273. 		//prlong longf("Node to be deleted not found\n");
  274. 		return;
  275. 	}
  276. 	NODEPTR tmp = Z->p;
  277. 	while (tmp != NILPTR) {
  278. 		tmp->size--;
  279. 		tmp = tmp->p;
  280. 	}
  281. 	NODEPTR y = Z;
  282. 	long long yoc = y->color;
  283. 	NODEPTR x;
  284. 	if (Z->left == NILPTR) {
  285. 		x = Z->right;
  286. 		rbtransplant(treeroot,Z,Z->right);
  287. 	}
  288. 	else if (Z->right == NILPTR) {
  289. 		x = Z->left;
  290. 		rbtransplant(treeroot,Z,Z->left);
  291. 	}
  292. 	else {
  293. 		y = minimum(Z->right);
  294. 		yoc = y->color;
  295. 		x = y->right;
  296. 		if (y->p == Z)
  297. 			x->p = y;
  298. 		else {
  299. 			rbtransplant(treeroot,y,y->right);
  300. 			y->right = Z->right;
  301. 			y->right->p = y;
  302. 		}
  303. 		rbtransplant(treeroot,Z,y);
  304. 		y->left = Z->left;
  305. 		y->left->p = y;
  306. 		y->color = Z->color;
  307. 	}
  308. 	if (yoc == BLACK)
  309. 		rbdeletefixup(treeroot,x);
  310. }
  311. NODEPTR kth(NODEPTR treeroot, long long K) {
  312. 	long long currrank = treeroot->left->size + 1;
  313. 	NODEPTR y = treeroot;
  314. 	while (y != NILPTR && currrank != K) {
  315. 		if (K < currrank)
  316. 			y = y->left;
  317. 		else {
  318. 			K = K - currrank;
  319. 			y = y->right;
  320. 		}
  321. 		if (y == NILPTR)
  322. 			return NILPTR;
  323. 		currrank = y->left->size + 1;
  324. 	}
  325. 	return y;
  326. }
  327.  
  328. long long cnt(NODEPTR treeroot, long long x) {
  329. 	long long ans = 0;
  330. 	NODEPTR y = treeroot;
  331. 	while (y != NILPTR) {
  332. 		if (y->key > x)
  333. 			y = y->left;
  334. 		else if (y->key < x) {
  335. 			ans += y->left->size + 1;
  336. 			y = y->right;
  337. 		}
  338. 		else 
  339. 			return ans + y->left->size;
  340. 	}
  341. 	return ans;
  342. }
  343.  
  344. main()
  345. {
  346. 	NIL.left = NIL.right = NIL.p = NILPTR;
  347. 	NIL.color = BLACK;
  348. 	NODEPTR tree = NILPTR;
  349. 	long long Q;
  350. 	long long x, k;
  351. 	NODEPTR temp;
  352. 	char c[2];
  353. 	scanf("%lld", &Q);
  354. 	getchar();
  355. 	while (Q--) {
  356. 		scanf("%s", c);
  357. 		switch (c[0]) {
  358. 			case 'I':
  359. 				scanf("%lld", &x);
  360. 				rbinsert(&tree, x);
  361. 				break;
  362. 			case 'D':
  363. 				scanf("%lld", &x);
  364. 				rbdelete(&tree, x);
  365. 				break;
  366. 			case 'K':
  367. 				scanf("%lld", &k);
  368. 				temp = kth(tree, k);
  369. 				if (temp != NILPTR)
  370. 					printf("%lld\n", temp->key);
  371. 				else
  372. 					printf("invalid\n");
  373. 				break;
  374. 			case 'C':
  375. 				scanf("%lld", &x);
  376. 				printf("%lld\n", cnt(tree,x));
  377. 				break;
  378. 			default:
  379. 				break;
  380. 		}
  381. 	}
  382. 	/*inorder(tree);
  383. 	putchar('\n');
  384. 	preorder(tree);
  385. 	putchar('\n');*/
  386. 	return 0;
  387. }
  388.  
Success #stdin #stdout 0s 2388KB
comments (?)
stdin
copy 
8
I -1
I -1
I 2
C 0
K 2
D -1
K 1
K 2
stdout
copy 
1
2
2
invalid
https://ideone.com/Z4sWHL
language:
C (gcc 8.3)
created:
10 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!