fork download
  1. #include <bits/stdc++.h>
  2. using namespace std;
  3.  
  4. // =====================================================================================
  5. // This file contains a collection of functions and a data structure for
  6. // MERGE SORT TREE.
  7. //
  8. // A Merge Sort Tree is a Segment Tree where each node stores a SORTED VECTOR
  9. // of the elements in its segment. It is used to answer range queries that ask
  10. // about order statistics (e.g., count of elements ≤ X) in a subarray [L, R].
  11. //
  12. // All functions are documented with:
  13. // - What they solve
  14. // - Input parameters
  15. // - Return value
  16. // - Time complexity
  17. // - Constraints / assumptions
  18. // - Important notes
  19. //
  20. // You can use these functions as black boxes. Read the comments carefully
  21. // to understand what each one does and when to use it.
  22. // =====================================================================================
  23.  
  24. // =====================================================================================
  25. // 1) MERGE SORT TREE CLASS
  26. // A class that builds a merge sort tree from a vector and provides query methods.
  27. // =====================================================================================
  28.  
  29. class MergeSortTree {
  30. private:
  31. int n; // size of the original array
  32. vector<vector<int>> tree; // tree[node] = sorted vector of that segment
  33.  
  34. // Build the tree recursively from the array a.
  35. // Parameters:
  36. // - node: current tree node index (1‑based)
  37. // - l, r: segment boundaries [l, r] (0‑based inclusive)
  38. // - a: original array
  39. void build(int node, int l, int r, const vector<int>& a) {
  40. if (l == r) {
  41. tree[node].push_back(a[l]);
  42. return;
  43. }
  44. int mid = (l + r) / 2;
  45. build(node * 2, l, mid, a);
  46. build(node * 2 + 1, mid + 1, r, a);
  47. // Merge the two sorted vectors from children
  48. merge(tree[node * 2].begin(), tree[node * 2].end(),
  49. tree[node * 2 + 1].begin(), tree[node * 2 + 1].end(),
  50. back_inserter(tree[node]));
  51. }
  52.  
  53. // Internal query: count elements ≤ X in range [ql, qr].
  54. int queryLessEqual(int node, int l, int r, int ql, int qr, int X) const {
  55. if (qr < l || r < ql) return 0; // no overlap
  56. if (ql <= l && r <= qr) { // full cover
  57. return upper_bound(tree[node].begin(), tree[node].end(), X) - tree[node].begin();
  58. }
  59. int mid = (l + r) / 2;
  60. return queryLessEqual(node * 2, l, mid, ql, qr, X) +
  61. queryLessEqual(node * 2 + 1, mid + 1, r, ql, qr, X);
  62. }
  63.  
  64. public:
  65. // Constructor: builds the tree from array 'a'.
  66. // Parameters:
  67. // - a: vector of integers (can be negative, zero, positive).
  68. // Time complexity: O(n log n), where n = a.size().
  69. // Memory: O(n log n) because each element appears in log n nodes.
  70. MergeSortTree(const vector<int>& a) {
  71. n = a.size();
  72. tree.resize(4 * n + 5);
  73. if (n > 0) build(1, 0, n - 1, a);
  74. }
  75.  
  76. // 1.1) Query: count of elements in range [L, R] that are <= X.
  77. // Parameters:
  78. // - L, R: inclusive indices of the subarray (0‑based).
  79. // - X: the upper bound value.
  80. // Returns:
  81. // - number of elements a[i] such that L <= i <= R and a[i] <= X.
  82. // Time complexity: O(log² n) (visits O(log n) nodes, each does a binary search).
  83. // Constraint: L <= R, 0 <= L,R < n.
  84. // Note: if X is very small, result may be 0; if X is very large, result = length.
  85. int queryLessEqual(int L, int R, int X) const {
  86. if (L > R || n == 0) return 0;
  87. return queryLessEqual(1, 0, n - 1, L, R, X);
  88. }
  89.  
  90. // 1.2) Query: count of elements in range [L, R] that are > X.
  91. // Parameters:
  92. // - L, R: inclusive indices.
  93. // - X: the threshold value.
  94. // Returns:
  95. // - number of elements > X in the subarray.
  96. // Time complexity: O(log² n).
  97. // Note: uses queryLessEqual to compute: total length - count(<= X).
  98. int queryGreater(int L, int R, int X) const {
  99. if (L > R || n == 0) return 0;
  100. int len = R - L + 1;
  101. return len - queryLessEqual(L, R, X);
  102. }
  103.  
  104. // 1.3) Query: count of elements in range [L, R] that are between LOW and HIGH
  105. // (inclusive, i.e., LOW <= a[i] <= HIGH).
  106. // Parameters:
  107. // - L, R: inclusive indices.
  108. // - LOW, HIGH: the lower and upper bounds (LOW <= HIGH).
  109. // Returns:
  110. // - number of elements in [L,R] with value in [LOW, HIGH].
  111. // Time complexity: O(log² n).
  112. // Note: uses queryLessEqual twice: count(<= HIGH) - count(< LOW) =>
  113. // count(<= HIGH) - count(<= LOW-1).
  114. int queryInRange(int L, int R, int LOW, int HIGH) const {
  115. if (L > R || LOW > HIGH || n == 0) return 0;
  116. return queryLessEqual(L, R, HIGH) - queryLessEqual(L, R, LOW - 1);
  117. }
  118.  
  119. // 1.4) Query: find the k‑th smallest element (1‑based) in the subarray [L, R].
  120. // Parameters:
  121. // - L, R: inclusive indices.
  122. // - k: 1‑based order (1 = smallest, length = largest).
  123. // Returns:
  124. // - the value of the k‑th smallest element in the subarray.
  125. // - If k is out of range, the behaviour is undefined (you must ensure 1 <= k <= len).
  126. // Time complexity: O(log³ n) typically (binary search over value range,
  127. // each check O(log² n)). For a value range of size up to 2e9, ~31 * log² n.
  128. // Constraint: the elements must be comparable (integers). The array values
  129. // must fit in int. The binary search assumes values are between -1e9 and 1e9.
  130. // If your values can be outside this range, adjust LOW and HIGH accordingly.
  131. // Note: This implementation uses binary search on the value domain.
  132. int queryKthSmallest(int L, int R, int k) const {
  133. if (L > R || n == 0) return 0;
  134. int len = R - L + 1;
  135. if (k < 1 || k > len) return 0; // optional safety
  136.  
  137. // Adjust these bounds if your values can be outside [-1e9, 1e9]
  138. int low = -1000000000, high = 1000000000;
  139. while (low < high) {
  140. int mid = low + (high - low) / 2;
  141. int cnt = queryLessEqual(L, R, mid);
  142. if (cnt >= k)
  143. high = mid;
  144. else
  145. low = mid + 1;
  146. }
  147. return low;
  148. }
  149.  
  150. // 1.5) Query: find the k‑th smallest using coordinate compression (if values are known).
  151. // This version assumes that all possible values are integers and we can compress them.
  152. // It is faster if we have a sorted list of all unique values.
  153. // Parameters:
  154. // - L, R: inclusive indices.
  155. // - k: 1‑based order.
  156. // - sortedVals: a sorted vector of all unique values that may appear.
  157. // Returns:
  158. // - the k‑th smallest value.
  159. // Time complexity: O(log² n * log m) where m = sortedVals.size().
  160. // Constraint: sortedVals must contain all values from the array.
  161. int queryKthSmallestCompressed(int L, int R, int k, const vector<int>& sortedVals) const {
  162. if (L > R || n == 0) return 0;
  163. int len = R - L + 1;
  164. if (k < 1 || k > len) return 0;
  165.  
  166. int low = 0, high = (int)sortedVals.size() - 1;
  167. while (low < high) {
  168. int mid = (low + high) / 2;
  169. int cnt = queryLessEqual(L, R, sortedVals[mid]);
  170. if (cnt >= k)
  171. high = mid;
  172. else
  173. low = mid + 1;
  174. }
  175. return sortedVals[low];
  176. }
  177.  
  178. // 1.6) Query: count of pairs (i,j) with L <= i < j <= R and a[i] + a[j] <= K.
  179. // This is a common trick used in ECPC/ACPC problems.
  180. // Parameters:
  181. // - L, R: inclusive range of indices.
  182. // - K: the sum limit.
  183. // Returns:
  184. // - number of pairs (i,j) within [L,R] with i<j and a[i]+a[j] <= K.
  185. // Time complexity: O(log² n) * (??) Actually this is not directly supported
  186. // by a simple merge sort tree. The correct way is to use a fenwick tree
  187. // of order statistics or a merge sort tree with two pointers on the fly.
  188. // However, we can implement a function that for each element counts
  189. // how many previous elements in the range satisfy the condition using
  190. // repeated queries. That would be O(len * log² n), which is too slow.
  191. // Instead, we will NOT include this as a black box function because it
  192. // is not efficient. Instead, we provide a note.
  193. // Note: For range pair counting, use a different approach (e.g., Mo's algorithm,
  194. // or Fenwick tree offline). Merge sort tree is not ideal for this.
  195. // This function is intentionally omitted.
  196. };
  197.  
  198. // =====================================================================================
  199. // 2) HELPER FUNCTIONS (not class methods) for common tasks using Merge Sort Tree
  200. // =====================================================================================
  201.  
  202. // 2.1) Count inversions in an array using merge sort (O(n log n)).
  203. // This is not a merge sort tree, but a classic divide-and-conquer.
  204. // Included here because it is often used in similar problems.
  205. // Parameters:
  206. // - a: vector of integers (will be modified during the process).
  207. // Returns:
  208. // - the number of inversions (i < j and a[i] > a[j]).
  209. // Time complexity: O(n log n).
  210. // Constraint: none, works for any integers.
  211. // Note: This function modifies the input (sorts it). If you need the original,
  212. // pass a copy.
  213. long long countInversionsMergeSort(vector<int>& a) {
  214. int n = a.size();
  215. if (n <= 1) return 0;
  216. int mid = n / 2;
  217. vector<int> left(a.begin(), a.begin() + mid);
  218. vector<int> right(a.begin() + mid, a.end());
  219. long long inv = countInversionsMergeSort(left) + countInversionsMergeSort(right);
  220. int i = 0, j = 0, k = 0;
  221. while (i < (int)left.size() && j < (int)right.size()) {
  222. if (left[i] <= right[j]) {
  223. a[k++] = left[i++];
  224. } else {
  225. a[k++] = right[j++];
  226. inv += (int)left.size() - i;
  227. }
  228. }
  229. while (i < (int)left.size()) a[k++] = left[i++];
  230. while (j < (int)right.size()) a[k++] = right[j++];
  231. return inv;
  232. }
  233.  
  234. // 2.2) Count the number of subarrays in [L, R] with sum <= K using prefix sums + Merge Sort Tree.
  235. // This is an advanced trick: for an array (can contain negatives?), we need a
  236. // different approach. But if all elements are non-negative, we can use sliding window.
  237. // For general values, we can compute prefix sums P[0..n], then count pairs (i,j)
  238. // with L <= i < j <= R and P[j] - P[i] <= K => P[j] <= P[i] + K.
  239. // This can be answered by a Merge Sort Tree built on prefix sums, but we need
  240. // to ensure index order. Actually we can iterate over j and query how many previous
  241. // prefix sums are >= P[j] - K using a Fenwick tree over compressed prefix sums.
  242. // So we provide a function that uses a Fenwick tree offline, not a merge sort tree.
  243. // For the purpose of this template, we will mention it but not implement it,
  244. // because it is not a pure merge sort tree application.
  245.  
  246. // =====================================================================================
  247. // 3) ADVANCED / TRICKS THAT APPEARED IN ECPC/ACPC CONTESTS
  248. // =====================================================================================
  249.  
  250. // (ASSUMPTION) The class MergeSortTree is already defined as in the previous template.
  251. // It provides:
  252. // - queryLessEqual(L, R, X) -> count of elements <= X in [L,R]
  253. // - queryInRange(L, R, LOW, HIGH) -> count of elements in [LOW, HIGH] in [L,R]
  254. // - queryKthSmallest(L, R, k) -> k-th smallest (1-indexed) in [L,R]
  255.  
  256. // =====================================================================================
  257. // 3.1) PROBLEM: Count how many elements in subarray [L, R] are ≤ X.
  258. // Parameters:
  259. // - mst: a MergeSortTree object built from the original array.
  260. // - L, R: inclusive 0-based indices of the subarray.
  261. // - X: the upper bound value (inclusive).
  262. // Returns:
  263. // - the number of elements a[i] with L ≤ i ≤ R and a[i] ≤ X.
  264. // Time complexity: O(log² N) (calls the tree's query method).
  265. // Constraint: L ≤ R and 0 ≤ L,R < array size.
  266. // Note: If X is very small, result may be 0.
  267. // =====================================================================================
  268. int countElementsLE(const MergeSortTree& mst, int L, int R, int X) {
  269. return mst.queryLessEqual(L, R, X);
  270. }
  271.  
  272. // =====================================================================================
  273. // 3.2) PROBLEM: Find the median of the subarray [L, R].
  274. // Parameters:
  275. // - mst: a MergeSortTree object.
  276. // - L, R: inclusive 0-based indices.
  277. // Returns:
  278. // - the median value. For odd length, it is the middle element.
  279. // For even length, it returns the upper median (the (len/2 + 1)-th smallest).
  280. // Time complexity: O(log³ N) (binary search inside the tree).
  281. // Constraint: L ≤ R and array is not empty.
  282. // Note: Median is defined as the element at position (len + 1) / 2
  283. // (1-indexed) which gives the upper median for even lengths.
  284. // =====================================================================================
  285. int subarrayMedian(const MergeSortTree& mst, int L, int R) {
  286. int len = R - L + 1;
  287. int k = (len + 1) / 2; // 1-indexed position of the median
  288. return mst.queryKthSmallest(L, R, k);
  289. }
  290.  
  291. // =====================================================================================
  292. // 3.3) PROBLEM: Count the number of contiguous subarrays whose sum is in the range [A, B].
  293. // Parameters:
  294. // - nums: vector of integers.
  295. // - A: lower bound of the sum (inclusive).
  296. // - B: upper bound of the sum (inclusive).
  297. // Returns:
  298. // - total count of subarrays with sum between A and B.
  299. // Time complexity: O(N) where N = nums.size().
  300. // CONSTRAINT: This function works ONLY if all elements in 'nums' are
  301. // NON-NEGATIVE. If negative numbers are present, the answer
  302. // will be WRONG because the sliding window technique relies
  303. // on the monotonicity of sums when extending the window.
  304. // Note: Uses the trick: count(≤ B) - count(≤ A-1).
  305. // The helper 'countSubarraysAtMost' uses two pointers.
  306. // =====================================================================================
  307. long long countSubarraysAtMost(const vector<int>& nums, int target) {
  308. if (target < 0) return 0; // sums are non-negative, so none can be ≤ negative
  309. int n = nums.size();
  310. long long ans = 0;
  311. int l = 0;
  312. long long sum = 0;
  313. for (int r = 0; r < n; r++) {
  314. sum += nums[r];
  315. while (sum > target) {
  316. sum -= nums[l++];
  317. }
  318. ans += (r - l + 1); // all subarrays ending at 'r' with start ≥ l are valid
  319. }
  320. return ans;
  321. }
  322.  
  323. long long countSubarraysSumInRange(const vector<int>& nums, int A, int B) {
  324. if (A > B) return 0;
  325. // count(≤ B) - count(≤ A-1)
  326. return countSubarraysAtMost(nums, B) - countSubarraysAtMost(nums, A - 1);
  327. }
  328.  
  329. // =====================================================================================
  330. // 3.4) PROBLEM: Count the number of pairs (i, j) with i < j and |a[i] - a[j]| ≤ K
  331. // over the WHOLE array.
  332. // Parameters:
  333. // - nums: vector of integers (will be modified/sorted internally).
  334. // - K: the maximum allowed absolute difference.
  335. // Returns:
  336. // - total number of valid pairs.
  337. // Time complexity: O(N log N) due to sorting, then O(N) two-pointer scan.
  338. // Constraint: none; works for positive and negative numbers.
  339. // Note: This is for the entire array. For range queries [L,R], a more
  340. // complex offline approach (e.g., Mo's algorithm with a Fenwick tree)
  341. // is required, which is NOT included in this template.
  342. // =====================================================================================
  343. long long countPairsDiffAtMostK(vector<int>& nums, int K) {
  344. sort(nums.begin(), nums.end());
  345. int n = nums.size();
  346. long long ans = 0;
  347. int r = 0;
  348. for (int l = 0; l < n; l++) {
  349. if (r < l) r = l;
  350. while (r + 1 < n && nums[r + 1] - nums[l] <= K) {
  351. r++;
  352. }
  353. ans += (r - l); // pairs (l, l+1) ... (l, r)
  354. }
  355. return ans;
  356. }
  357.  
  358. // =====================================================================================
  359. // 3.5) PROBLEM: Count the number of elements in [L, R] that are between LOW and HIGH
  360. // (inclusive, i.e., LOW ≤ a[i] ≤ HIGH).
  361. // Parameters:
  362. // - mst: a MergeSortTree object.
  363. // - L, R: inclusive 0-based indices.
  364. // - LOW, HIGH: inclusive lower and upper bounds.
  365. // Returns:
  366. // - count of elements in the subarray with value in [LOW, HIGH].
  367. // Time complexity: O(log² N).
  368. // Constraint: LOW ≤ HIGH.
  369. // Note: This is essentially a wrapper for 'queryInRange'.
  370. // =====================================================================================
  371. int countElementsInRange(const MergeSortTree& mst, int L, int R, int LOW, int HIGH) {
  372. return mst.queryInRange(L, R, LOW, HIGH);
  373. }
  374.  
  375. // =====================================================================================
  376. // 3.6) PROBLEM: Find the K-th smallest element in the subarray [L, R].
  377. // Parameters:
  378. // - mst: a MergeSortTree object.
  379. // - L, R: inclusive 0-based indices.
  380. // - K: the order (1-indexed). For example, K=1 returns the minimum,
  381. // K=length returns the maximum.
  382. // Returns:
  383. // - the value of the K-th smallest element.
  384. // Time complexity: O(log³ N).
  385. // Constraint: 1 ≤ K ≤ (R - L + 1).
  386. // Note: This is a direct wrapper for 'queryKthSmallest'.
  387. // =====================================================================================
  388. int kthSmallestInSubarray(const MergeSortTree& mst, int L, int R, int K) {
  389. return mst.queryKthSmallest(L, R, K);
  390. }
  391.  
  392. // =====================================================================================
  393. // 3.7) PROBLEM: Count the number of elements in subarray [L, R] that are STRICTLY
  394. // less than X (i.e., a[i] < X).
  395. // Parameters:
  396. // - mst: a MergeSortTree object.
  397. // - L, R: inclusive 0-based indices.
  398. // - X: the threshold (exclusive).
  399. // Returns:
  400. // - count of elements < X in the subarray.
  401. // Time complexity: O(log² N).
  402. // Note: Since queryLessEqual counts elements ≤ X, we simply pass X-1.
  403. // For floating point values, this trick doesn't work; but this template
  404. // assumes integer values.
  405. // =====================================================================================
  406. int countStrictlyLess(const MergeSortTree& mst, int L, int R, int X) {
  407. return mst.queryLessEqual(L, R, X - 1);
  408. }
  409.  
  410. // =====================================================================================
  411. // 3.8) PROBLEM: Update a position (point update) in the array and still be able
  412. // to answer merge-sort-tree style queries (like count ≤ X).
  413. //
  414. // ------- IMPORTANT WARNING -------
  415. // The classic Merge Sort Tree (as defined in the previous template) does NOT
  416. // support efficient point updates. If you change one element, you would have
  417. // to rebuild the sorted vectors for all nodes on the path from the leaf to
  418. // the root. Rebuilding one node costs O(size of the node). In the worst case,
  419. // a single point update costs O(N log N), which is too slow for most problems.
  420. //
  421. // If your problem requires updates, use one of these alternatives:
  422. // 1) Fenwick Tree of Fenwick Trees (Fenwick Tree of Order Statistics):
  423. // - Supports point updates and range queries in O(log² N).
  424. // - Requires coordinate compression of all values (offline).
  425. // 2) Segment Tree of Balanced BSTs (e.g., std::multiset):
  426. // - Update: O(log² N), Query: O(log² N).
  427. // - Heavier constant factor.
  428. // 3) Sqrt Decomposition (Block decomposition):
  429. // - Simpler to implement, O(sqrt(N) * log(sqrt(N))) per query/update.
  430. //
  431. // The following function is a STUB to remind you that it is not supported.
  432. // If you call it, it will do nothing or return an error (assert false).
  433. // =====================================================================================
  434. void pointUpdate(MergeSortTree& mst, int pos, int newVal) {
  435. // This function intentionally does nothing.
  436. // The classic Merge Sort Tree does not support efficient point updates.
  437. // Rebuilding the whole tree or even a single node is O(N log N) in the
  438. // worst case, which defeats the purpose.
  439. // Please use a Fenwick tree of order statistics or a different data structure.
  440. // Uncomment the line below to cause a runtime error if accidentally called.
  441. // assert(false && "Point updates are not supported by the classic Merge Sort Tree.");
  442.  
  443. // If you absolutely must use this with rebuild, here is the inefficient way:
  444. // (DO NOT USE in contests unless N and Q are very small).
  445. // 1. Update the original array.
  446. // 2. Rebuild the entire tree: mst = MergeSortTree(updatedArray); // O(N log N)
  447. }
  448.  
  449. // =====================================================================================
  450. // 4) TEMPLATE FOR USING MERGE SORT TREE IN A COMPETITIVE PROGRAMMING SETTING
  451. // (Example main function – you can ignore or adapt)
  452. // =====================================================================================
  453.  
  454. int main() {
  455. ios::sync_with_stdio(false);
  456. cin.tie(nullptr);
  457.  
  458. // Example 1: Build a merge sort tree from an array.
  459. vector<int> arr = {5, 2, 8, 1, 9, 3, 7, 4, 6};
  460. MergeSortTree mst(arr);
  461.  
  462. // Query: count elements <= 5 in subarray [2, 6] (0‑based)
  463. int L = 2, R = 6, X = 5;
  464. cout << "Number <= 5 in [2,6]: " << mst.queryLessEqual(L, R, X) << '\n'; // Expected: 3 (elements 1,3,4)
  465.  
  466. // Query: count elements in [3, 7] in subarray [1,5]
  467. cout << "Elements in [3,7] in [1,5]: " << mst.queryInRange(1, 5, 3, 7) << '\n';
  468.  
  469. // Query: 3rd smallest in whole array
  470. cout << "3rd smallest in whole array: " << mst.queryKthSmallest(0, 8, 3) << '\n';
  471.  
  472. // Example 2: Compressed kth smallest
  473. vector<int> vals = {1, 2, 3, 4, 5, 6, 7, 8, 9}; // all unique values
  474. cout << "3rd smallest (compressed): " << mst.queryKthSmallestCompressed(0, 8, 3, vals) << '\n';
  475.  
  476. return 0;
  477. }
  478.  
  479. // =====================================================================================
  480. // 5) GLOSSARY OF TERMS USED
  481. // =====================================================================================
  482. //
  483. // - Merge Sort Tree: A segment tree where each node stores a sorted list of the elements
  484. // in its segment. It is built by merging the sorted lists of children (like merge sort).
  485. // - Node: a node in the segment tree, representing a contiguous segment of the array.
  486. // - Segment: a contiguous subarray of the original array.
  487. // - Order statistic: a value that describes the position of an element when the data is sorted
  488. // (e.g., the smallest, the 10th smallest).
  489. // - K‑th smallest: the element that would be at position k if the subarray were sorted
  490. // in ascending order (1‑based).
  491. // - Binary search: a method to find a target value in a sorted array by repeatedly halving
  492. // the search space. Here used to find the k‑th smallest.
  493. // - Logarithmic: O(log n) time complexity, meaning the time grows slowly as n increases.
  494. // - Coordinate compression: mapping each distinct value to a smaller integer (its rank)
  495. // to use in data structures like Fenwick trees; here used to speed up kth smallest.
  496. // - Fenwick tree (BIT): a different data structure for prefix sums and order statistics;
  497. // not part of this template, but mentioned for context.
  498. // - Mo's algorithm: an offline algorithm for range queries, not used here.
  499. // - Inversion: a pair (i, j) such that i < j and a[i] > a[j].
  500. // =====================================================================================
Success #stdin #stdout 0.01s 5304KB
stdin
Standard input is empty
stdout
Number <= 5 in [2,6]: 2
Elements in [3,7] in [1,5]: 1
3rd smallest in whole array: 3
3rd smallest (compressed): 3