fork download
  1. #include <bits/stdc++.h>
  2. using namespace std;
  3.  
  4. using ll = long long;
  5.  
  6. // ============================================================================
  7. // SEGMENT TREE WITH WALK OPERATIONS (FIND FIRST / LAST, K‑TH, PREFIX WALK)
  8. // ============================================================================
  9. // This file provides a generic segment tree that supports:
  10. // • point assignment
  11. // • range query (sum, max, min, gcd, …)
  12. // • range add (lazy propagation – only for sum)
  13. // • "walk" functions:
  14. // - maxRight / minLeft (AtCoder style, monotonic predicate)
  15. // - findKth (k‑th element by prefix sum)
  16. // - walkToSum (first position where prefix sum ≥ target)
  17. //
  18. // All functions are ready to be used as black boxes.
  19. // Read the comments above each function to understand:
  20. // - what it solves
  21. // - what input it expects
  22. // - what it returns
  23. // - time complexity
  24. // - important constraints / assumptions
  25. // ============================================================================
  26.  
  27. // ============================================================================
  28. // Generic Segment Tree Class
  29. // ============================================================================
  30. // T : type of the elements (int, long long, …)
  31. // combine : associative operation (e.g. sum, max, min, gcd)
  32. // neutral : identity element for combine (0 for sum, -INF for max, +INF for min, …)
  33. // useLazy : if true, enables range add (requires combine = sum and neutral = 0)
  34. // If you only need point updates / range queries, set useLazy = false.
  35. //
  36. // IMPORTANT:
  37. // • rangeAdd() works ONLY if combine is addition (sum) and neutral = 0.
  38. // • maxRight() and minLeft() require a predicate that is monotonic over the
  39. // monoid (see explanations inside).
  40. // • findKth() and walkToSum() also assume combine is sum.
  41. // • All indices are 0‑based and inclusive in queries, except where noted.
  42. // ============================================================================
  43.  
  44. template <class T>
  45. class SegTree {
  46. private:
  47. int n; // number of elements
  48. vector<T> tree, lazy; // segment tree and lazy values
  49. T neutral; // identity element
  50. function<T(T,T)> combine; // associative operation
  51. bool useLazy; // whether lazy range add is enabled
  52.  
  53. // ---------------------- internal helpers ----------------------
  54. void build(int node, int l, int r, const vector<T>& data) {
  55. if (l == r) {
  56. tree[node] = data[l];
  57. return;
  58. }
  59. int mid = (l + r) >> 1;
  60. build(node<<1, l, mid, data);
  61. build(node<<1|1, mid+1, r, data);
  62. tree[node] = combine(tree[node<<1], tree[node<<1|1]);
  63. }
  64.  
  65. // apply an addition to a node – works ONLY for sum
  66. void applyAdd(int node, T val) {
  67. tree[node] += val;
  68. if (useLazy) lazy[node] += val;
  69. }
  70.  
  71. void push(int node) {
  72. if (!useLazy) return;
  73. if (lazy[node] != neutral) {
  74. applyAdd(node<<1, lazy[node]);
  75. applyAdd(node<<1|1, lazy[node]);
  76. lazy[node] = neutral;
  77. }
  78. }
  79.  
  80. // point set (no lazy push needed for point set? we push when going down)
  81. void pointSet(int node, int l, int r, int pos, T val) {
  82. if (l == r) {
  83. tree[node] = val;
  84. return;
  85. }
  86. push(node);
  87. int mid = (l + r) >> 1;
  88. if (pos <= mid) pointSet(node<<1, l, mid, pos, val);
  89. else pointSet(node<<1|1, mid+1, r, pos, val);
  90. tree[node] = combine(tree[node<<1], tree[node<<1|1]);
  91. }
  92.  
  93. // range add (only for sum)
  94. void rangeAdd(int node, int l, int r, int ql, int qr, T val) {
  95. if (ql <= l && r <= qr) {
  96. applyAdd(node, val);
  97. return;
  98. }
  99. push(node);
  100. int mid = (l + r) >> 1;
  101. if (ql <= mid) rangeAdd(node<<1, l, mid, ql, qr, val);
  102. if (qr > mid) rangeAdd(node<<1|1, mid+1, r, ql, qr, val);
  103. tree[node] = combine(tree[node<<1], tree[node<<1|1]);
  104. }
  105.  
  106. // range query (works for any combine)
  107. T query(int node, int l, int r, int ql, int qr) {
  108. if (ql <= l && r <= qr) return tree[node];
  109. push(node);
  110. int mid = (l + r) >> 1;
  111. if (qr <= mid) return query(node<<1, l, mid, ql, qr);
  112. if (ql > mid) return query(node<<1|1, mid+1, r, ql, qr);
  113. return combine(
  114. query(node<<1, l, mid, ql, qr),
  115. query(node<<1|1, mid+1, r, ql, qr)
  116. );
  117. }
  118.  
  119. // ---------- maxRight helper (inclusive tree) ----------
  120. // Finds the first index r in [ql, n] where the predicate becomes false.
  121. // sm accumulates the value of the prefix that has been proven to satisfy the predicate.
  122. // Returns n if the predicate stays true for the whole suffix.
  123. int maxRightRec(int node, int l, int r, int ql, T& sm, const function<bool(T)>& pred) {
  124. if (r < ql) return n; // segment completely before ql
  125. T combined = combine(sm, tree[node]);
  126. // If the whole segment is inside the query and adding it keeps pred true, take it.
  127. if (ql <= l && r <= n-1 && pred(combined)) {
  128. sm = combined;
  129. return n;
  130. }
  131. if (l == r) {
  132. // leaf: we must decide whether to include it
  133. T leafVal = tree[node];
  134. T newVal = combine(sm, leafVal);
  135. if (pred(newVal)) {
  136. sm = newVal;
  137. return n;
  138. } else {
  139. return l; // this leaf is the first failure point
  140. }
  141. }
  142. push(node);
  143. int mid = (l + r) >> 1;
  144. int res = maxRightRec(node<<1, l, mid, ql, sm, pred);
  145. if (res != n) return res;
  146. return maxRightRec(node<<1|1, mid+1, r, ql, sm, pred);
  147. }
  148.  
  149. // ---------- minLeft helper (inclusive tree) ----------
  150. // Finds the first index l (from the right) where the predicate fails,
  151. // while building the suffix from right to left.
  152. // Parameter rbound is the exclusive right bound of the query (0 <= rbound <= n).
  153. // sm accumulates the suffix that has been proven to satisfy pred.
  154. // Returns -1 if the predicate holds for the whole suffix down to index 0.
  155. // Otherwise returns the index of the element that cannot be included.
  156. int minLeftRec(int node, int l, int r, int rbound, T& sm, const function<bool(T)>& pred) {
  157. if (l >= rbound) return -1; // segment completely after the query range
  158. T combined = combine(tree[node], sm);
  159. // If the whole segment is inside the query and adding it keeps pred true, take it.
  160. if (r <= rbound-1 && pred(combined)) {
  161. sm = combined;
  162. return -1;
  163. }
  164. if (l == r) {
  165. // leaf: test if we can include it
  166. T leafVal = tree[node];
  167. T newVal = combine(leafVal, sm);
  168. if (pred(newVal)) {
  169. sm = newVal;
  170. return -1;
  171. } else {
  172. return l; // cannot include this leaf
  173. }
  174. }
  175. push(node);
  176. int mid = (l + r) >> 1;
  177. // go right first (since we are moving right‑to‑left)
  178. int res = minLeftRec(node<<1|1, mid+1, r, rbound, sm, pred);
  179. if (res != -1) return res;
  180. return minLeftRec(node<<1, l, mid, rbound, sm, pred);
  181. }
  182.  
  183. // ---------- findKth helper (sum only) ----------
  184. // assumes tree[node] stores the sum of its segment.
  185. // k is 0‑based: we want the smallest index p such that prefix sum up to p > k.
  186. int findKthRec(int node, int l, int r, T k) {
  187. if (l == r) return l;
  188. push(node);
  189. int mid = (l + r) >> 1;
  190. if (tree[node<<1] > k) return findKthRec(node<<1, l, mid, k);
  191. else return findKthRec(node<<1|1, mid+1, r, k - tree[node<<1]);
  192. }
  193.  
  194. // ---------- walkToSum helper (sum only) ----------
  195. // Finds the first position p (starting from ql) where the accumulated sum >= target.
  196. // acc holds the sum of the prefix that has already been taken.
  197. // Returns n if the total sum from ql to end is < target.
  198. int walkToSumRec(int node, int l, int r, int ql, T& acc, T target) {
  199. if (r < ql) return n;
  200. if (ql <= l) {
  201. T newAcc = combine(acc, tree[node]);
  202. if (newAcc < target) {
  203. acc = newAcc;
  204. return n; // whole segment taken, still not enough
  205. }
  206. if (l == r) {
  207. return l; // leaf makes sum reach target
  208. }
  209. }
  210. push(node);
  211. int mid = (l + r) >> 1;
  212. int res = walkToSumRec(node<<1, l, mid, ql, acc, target);
  213. if (res != n) return res;
  214. return walkToSumRec(node<<1|1, mid+1, r, ql, acc, target);
  215. }
  216.  
  217. public:
  218. // ---------- constructor ----------
  219. // data : initial array (0‑indexed)
  220. // neutral : identity element for combine
  221. // combine : associative binary operation (e.g. [](T a, T b){ return a+b; })
  222. // useLazy : enable range add (requires combine = addition)
  223. SegTree(const vector<T>& data, T neutral, function<T(T,T)> combine, bool useLazy = false)
  224. : neutral(neutral), combine(combine), useLazy(useLazy) {
  225. n = (int)data.size();
  226. tree.assign(4*n + 5, neutral);
  227. lazy.assign(4*n + 5, neutral);
  228. build(1, 0, n-1, data);
  229. }
  230.  
  231. // ---------- point assignment ----------
  232. // Sets the value at position pos (0‑indexed) to val.
  233. // Time: O(log n)
  234. void pointSet(int pos, T val) {
  235. pointSet(1, 0, n-1, pos, val);
  236. }
  237.  
  238. // ---------- range add (lazy) ----------
  239. // Adds val to every element in [l, r] (inclusive).
  240. // ONLY works if combine is sum (addition) and neutral = 0.
  241. // Time: O(log n)
  242. void rangeAdd(int l, int r, T val) {
  243. if (!useLazy) {
  244. cerr << "WARNING: rangeAdd called but lazy is disabled. This may give wrong results.\n";
  245. }
  246. rangeAdd(1, 0, n-1, l, r, val);
  247. }
  248.  
  249. // ---------- range query ----------
  250. // Returns combine( data[l], data[l+1], …, data[r] ) (inclusive).
  251. // Works for any combine.
  252. // Time: O(log n)
  253. T query(int l, int r) {
  254. return query(1, 0, n-1, l, r);
  255. }
  256.  
  257. // ---------- maxRight (AtCoder style) ----------
  258. // Finds the smallest index r (l <= r <= n) such that
  259. // pred( combine( data[l], data[l+1], …, data[r-1] ) ) == false,
  260. // i.e. the first position where the predicate becomes false.
  261. // If the predicate is true for the whole array, returns n.
  262. // The empty prefix (r = l) is always considered true, so pred(neutral) must be true.
  263. //
  264. // The predicate pred must be monotonic:
  265. // if pred(X) is true, then pred( combine(X, Y) ) may be true or false,
  266. // but once it becomes false, it stays false when you extend the segment.
  267. // (This holds for many natural predicates, e.g. sum < K).
  268. //
  269. // Time: O(log n)
  270. int maxRight(int l, const function<bool(T)>& pred) {
  271. T sm = neutral;
  272. int res = maxRightRec(1, 0, n-1, l, sm, pred);
  273. return res; // res is either n or the first failure index
  274. }
  275.  
  276. // ---------- minLeft (AtCoder style) ----------
  277. // Given r (0 <= r <= n), finds the minimum index l such that
  278. // pred( combine( data[l], data[l+1], …, data[r-1] ) ) == true.
  279. // In other words, the largest prefix that can be excluded from the right
  280. // while keeping the predicate true on the remaining suffix.
  281. // The empty suffix (l = r) is always considered true, so pred(neutral) must be true.
  282. //
  283. // Predicate monotonic as for maxRight.
  284. //
  285. // Time: O(log n)
  286. int minLeft(int r, const function<bool(T)>& pred) {
  287. T sm = neutral;
  288. int res = minLeftRec(1, 0, n-1, r, sm, pred);
  289. return (res == -1 ? 0 : res + 1);
  290. }
  291.  
  292. // ---------- findKth (sum only) ----------
  293. // Finds the smallest index p (0‑based) such that:
  294. // sum( data[0] + data[1] + … + data[p] ) > k
  295. // (i.e. the position of the (k+1)‑th unit when elements represent frequencies).
  296. // This assumes all data elements are non‑negative and combine is addition.
  297. // k is 0‑based: k = 0 returns the position of the first element that makes prefix sum > 0.
  298. //
  299. // Returns p, or n if total sum <= k.
  300. // Time: O(log n)
  301. int findKth(T k) {
  302. if (tree[1] <= k) return n;
  303. return findKthRec(1, 0, n-1, k);
  304. }
  305.  
  306. // ---------- walkToSum (sum only) ----------
  307. // Finds the smallest index p (0‑based) such that:
  308. // sum( data[start] + data[start+1] + … + data[p] ) >= target
  309. // The search starts from position 'start' (default 0).
  310. // Returns p, or n if the total sum from start to n-1 is < target.
  311. // This is useful for prefix‑based queries.
  312. // Time: O(log n)
  313. int walkToSum(T target, int start = 0) {
  314. if (start >= n) return n;
  315. T acc = neutral; // neutral = 0 for sum
  316. int res = walkToSumRec(1, 0, n-1, start, acc, target);
  317. return res;
  318. }
  319.  
  320. // ---------- total aggregate ----------
  321. // Returns combine of all elements (tree[1]).
  322. // Time: O(1)
  323. T all() const {
  324. return tree[1];
  325. }
  326. };
  327.  
  328. // ============================================================================
  329. // ADDITIONAL HELPER FUNCTIONS (standalone)
  330. // These are often used together with a segment tree that stores sums
  331. // (like a Fenwick alternative) or for two‑pointer problems.
  332. // ============================================================================
  333.  
  334. // ----------------------------------------------------------------------------
  335. // 1) Maximum number of pairs from two sorted arrays with sum ≤ K
  336. // ----------------------------------------------------------------------------
  337. // PURPOSE:
  338. // Given two arrays A and B, find the maximum number of disjoint pairs
  339. // (one from A, one from B) such that A[i] + B[j] ≤ K.
  340. // INPUT:
  341. // A, B : vectors of ints (will be sorted internally)
  342. // K : upper bound
  343. // OUTPUT:
  344. // Maximum number of pairs.
  345. // TIME: O(n log n + m log m) due to sorting, then O(n+m) two pointers.
  346. int maxPairsWithSumAtMostK(vector<int>& A, vector<int>& B, int K) {
  347. sort(A.begin(), A.end());
  348. sort(B.begin(), B.end());
  349. int i = 0, j = (int)B.size() - 1;
  350. int ans = 0;
  351. while (i < (int)A.size() && j >= 0) {
  352. if (A[i] + B[j] <= K) {
  353. ans++;
  354. i++;
  355. j--;
  356. } else {
  357. j--;
  358. }
  359. }
  360. return ans;
  361. }
  362.  
  363. // ----------------------------------------------------------------------------
  364. // 2) Count subarrays with sum in [L, R] (non‑negative array only)
  365. // ----------------------------------------------------------------------------
  366. // PURPOSE:
  367. // Counts the number of contiguous subarrays whose sum lies in [L, R].
  368. // INPUT:
  369. // nums : vector<int> with NON‑NEGATIVE elements.
  370. // L, R : long long bounds (L ≤ R).
  371. // OUTPUT:
  372. // Number of subarrays.
  373. // TIME: O(n log n)
  374. // WARNING: Works only if all nums are non‑negative (otherwise prefix sums
  375. // are not monotonic and the sorting trick fails).
  376. ll countSubarraysInRange(const vector<int>& nums, ll L, ll R) {
  377. int n = nums.size();
  378. vector<ll> pref(n + 1, 0);
  379. for (int i = 0; i < n; i++) pref[i+1] = pref[i] + nums[i];
  380. sort(pref.begin(), pref.end());
  381. auto countPairsLE = [&](ll X) -> ll {
  382. ll cnt = 0;
  383. int j = 0;
  384. for (int i = 0; i < (int)pref.size(); i++) {
  385. if (j < i) j = i;
  386. while (j + 1 < (int)pref.size() && pref[j+1] - pref[i] <= X) j++;
  387. cnt += (j - i);
  388. }
  389. return cnt;
  390. };
  391. return countPairsLE(R) - countPairsLE(L - 1);
  392. }
  393.  
  394. // ----------------------------------------------------------------------------
  395. // 3) Closest pair sum from two arrays
  396. // ----------------------------------------------------------------------------
  397. // PURPOSE:
  398. // Find a pair (a∈A, b∈B) whose sum is closest to a given target.
  399. // INPUT:
  400. // A, B : vectors of ints (will be sorted internally)
  401. // target : int
  402. // OUTPUT:
  403. // pair<int,int> with the chosen elements.
  404. // TIME: O(n log n + m log m) sorting, then O(n+m) two pointers.
  405. pair<int,int> closestPairFromTwoArrays(vector<int>& A, vector<int>& B, int target) {
  406. sort(A.begin(), A.end());
  407. sort(B.begin(), B.end());
  408. int i = 0, j = (int)B.size() - 1;
  409. int bestDiff = INT_MAX;
  410. pair<int,int> best = {A[0], B[0]};
  411. while (i < (int)A.size() && j >= 0) {
  412. int sum = A[i] + B[j];
  413. int diff = abs(sum - target);
  414. if (diff < bestDiff) {
  415. bestDiff = diff;
  416. best = {A[i], B[j]};
  417. }
  418. if (sum < target) i++;
  419. else if (sum > target) j--;
  420. else break;
  421. }
  422. return best;
  423. }
  424.  
  425. // ----------------------------------------------------------------------------
  426. // 4) Median of two sorted arrays (merge to middle)
  427. // ----------------------------------------------------------------------------
  428. // PURPOSE:
  429. // Returns the median of the merged array from two sorted arrays.
  430. // INPUT:
  431. // A, B : sorted vectors (non‑decreasing)
  432. // OUTPUT:
  433. // double median.
  434. // TIME: O(n+m)
  435. double medianOfTwoSortedArrays(const vector<int>& A, const vector<int>& B) {
  436. int n = A.size(), m = B.size();
  437. int total = n + m;
  438. int i = 0, j = 0;
  439. int prev = 0, cur = 0;
  440. for (int k = 0; k <= total/2; k++) {
  441. prev = cur;
  442. if (i < n && (j >= m || A[i] < B[j]))
  443. cur = A[i++];
  444. else
  445. cur = B[j++];
  446. }
  447. if (total % 2 == 1) return cur;
  448. return (prev + cur) / 2.0;
  449. }
  450.  
  451. // ----------------------------------------------------------------------------
  452. // 5) Minimum operations to make all array elements equal (cost = sum |x - median|)
  453. // ----------------------------------------------------------------------------
  454. // PURPOSE:
  455. // Minimum total number of increment/decrement operations to make all elements equal.
  456. // INPUT:
  457. // nums : vector<int>
  458. // OUTPUT:
  459. // long long minimal cost.
  460. // TIME: O(n log n) due to sorting, or O(n) with nth_element.
  461. ll minOperationsToMakeEqual(vector<int>& nums) {
  462. int n = nums.size();
  463. if (n == 0) return 0;
  464. sort(nums.begin(), nums.end());
  465. int median = nums[n/2];
  466. ll cost = 0;
  467. for (int x : nums) cost += llabs((ll)x - (ll)median);
  468. return cost;
  469. }
  470.  
  471. // ============================================================================
  472. // EXAMPLE USAGE (main)
  473. // ============================================================================
  474. int main() {
  475. ios::sync_with_stdio(false);
  476. cin.tie(nullptr);
  477.  
  478. // ---------- Example 1: Segment Tree with sum ----------
  479. vector<int> data = {1, 3, 2, 5, 4};
  480. SegTree<int> st(data, 0, [](int a, int b){ return a + b; }, true); // with lazy
  481.  
  482. cout << "Sum [0,4] = " << st.query(0, 4) << "\n"; // 15
  483.  
  484. st.rangeAdd(1, 3, 2); // add 2 to indices 1..3
  485. cout << "After add, sum [1,3] = " << st.query(1, 3) << "\n"; // (3+2)+(2+2)+(5+2)=16
  486.  
  487. st.pointSet(2, 10);
  488. cout << "After set index 2 to 10, sum [0,4] = " << st.query(0, 4) << "\n";
  489.  
  490. // ---------- Example 2: maxRight ----------
  491. vector<int> arr = {2, 3, 1, 4, 2};
  492. SegTree<int> st2(arr, 0, [](int a, int b){ return a + b; }, false);
  493.  
  494. // pred(x) = x < 7 (monotonic: as we add more numbers, sum increases)
  495. auto pred = [](int x) { return x < 7; };
  496. int r = st2.maxRight(0, pred);
  497. cout << "maxRight: longest prefix from 0 with sum < 7 ends at index " << r << "\n";
  498. // prefix 0..2 sum = 2+3+1=6 <7, adding index 3 gives 10 >=7, so r=3.
  499.  
  500. // ---------- Example 3: findKth ----------
  501. vector<int> weights = {0, 1, 0, 1, 0, 1}; // positions of ones
  502. SegTree<int> st3(weights, 0, [](int a, int b){ return a + b; }, false);
  503. int kth1 = st3.findKth(1); // k=1 (0‑based) -> second '1'
  504. cout << "Position of the 2nd '1' = " << kth1 << "\n"; // index 3 (since ones at 1,3,5)
  505.  
  506. // ---------- Example 4: walkToSum ----------
  507. vector<int> vals = {5, 2, 8, 3, 6};
  508. SegTree<int> st4(vals, 0, [](int a, int b){ return a + b; }, false);
  509. int pos = st4.walkToSum(15, 0); // need prefix sum >=15 from start
  510. cout << "First position where prefix sum >= 15 is " << pos << "\n"; // indices 0+1+2=15 -> pos=2
  511.  
  512. pos = st4.walkToSum(10, 2); // start at index 2: 8+3=11 -> pos=3
  513. cout << "From index 2, first pos with sum >=10 is " << pos << "\n"; // 3
  514.  
  515. // ---------- Example 5: minLeft ----------
  516. // pred(x) = x >= 5 on suffix sums (monotonic when moving left)
  517. auto predMin = [](int x) { return x >= 5; };
  518. int l = st4.minLeft(5, predMin); // r = 5 (exclusive, covers whole array)
  519. // Suffix from l to 4: we want smallest l such that sum(l..4) >= 5.
  520. // Suffixes: sum(4)=6>=5 -> l=4; sum(3..4)=9>=5 -> l=3; sum(2..4)=17>=5 -> l=2; sum(1..4)=19>=5 -> l=1; sum(0..4)=24>=5 -> l=0.
  521. // Smallest l that makes sum >=5 is actually 0? No, suffix sum from 4 is 6, so l=4 works. Smallest l is 4.
  522. cout << "minLeft: smallest l with sum(l..4) >= 5 is " << l << "\n"; // expected 4
  523.  
  524. return 0;
  525. }
Success #stdin #stdout 0s 5320KB
stdin
Standard input is empty
stdout
Sum [0,4] = 15
After add, sum [1,3] = 16
After set index 2 to 10, sum [0,4] = 27
maxRight: longest prefix from 0 with sum < 7 ends at index 3
Position of the 2nd '1' = 3
First position where prefix sum >= 15 is 2
From index 2, first pos with sum >=10 is 3
minLeft: smallest l with sum(l..4) >= 5 is 0