2.4.12 Merge Sort Tree

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106

/*

A merge sort tree is a static segment tree in which every node stores the sorted multiset of the
values in its range. The nodes are exactly the ranges visited by a merge sort, and each node's
sorted list is the merge of its two children's lists. This supports queries that count, within a
range of indices, how many values fall below a threshold or inside a value interval: decompose the
index range into O(log n) canonical nodes and binary search each node's sorted list.

The structure is built once and not modified afterward. All index ranges are inclusive
$[`lo`, `hi`]$ with 0-based indices, and values may be of any comparable type.

- `MergeSortTree<T>(a)` builds the tree over the array `a`.
- `size()` returns the size of the array.
- `count_leq(lo, hi, x)` returns the number of indices `i` in $[`lo`, `hi`]$ with `a[i]` $\leq$ `x`.
- `count_in(lo, hi, x, y)` returns the number of indices `i` $\in [`lo`, `hi`]$ such that `a[i]`
   $\in [`x`, `y`]$.

Time Complexity:
- O(n log n) per call to the constructor, where $n$ is the size of the array.
- O(1) per call to `size()`.
- O(log^2 n) per call to `count_leq()` and `count_in()`.

Space Complexity:
- O(n log n) for storage of the tree.
- O(log n) auxiliary stack space per query.

*/

#include <algorithm>
#include <cassert>
#include <iterator>
#include <vector>

template<typename T>
class MergeSortTree {
  int len;
  std::vector<std::vector<T>> tree;

  void build(int i, int lo, int hi, const std::vector<T> &a) {
    if (lo == hi) {
      tree[i] = {a[lo]};
      return;
    }
    int mid = lo + (hi - lo) / 2;
    build(i * 2, lo, mid, a);
    build(i * 2 + 1, mid + 1, hi, a);
    std::merge(
        tree[i * 2].begin(), tree[i * 2].end(), tree[i * 2 + 1].begin(), tree[i * 2 + 1].end(),
        std::back_inserter(tree[i])
    );
  }

  template<typename Fn>
  int query(int i, int lo, int hi, int tgt_lo, int tgt_hi, const Fn &count_node) const {
    if (tgt_hi < lo || hi < tgt_lo) {
      return 0;
    }
    if (tgt_lo <= lo && hi <= tgt_hi) {
      return count_node(tree[i]);
    }
    int mid = lo + (hi - lo) / 2;
    return query(i * 2, lo, mid, tgt_lo, tgt_hi, count_node) +
           query(i * 2 + 1, mid + 1, hi, tgt_lo, tgt_hi, count_node);
  }

 public:
  explicit MergeSortTree(const std::vector<T> &a) : len(a.size()), tree(4 * a.size()) {
    assert(len > 0);
    build(1, 0, len - 1, a);
  }

  int size() const { return len; }

  int count_leq(int lo, int hi, const T &x) const {
    assert(0 <= lo && lo <= hi && hi < len);
    return query(1, 0, len - 1, lo, hi, [&](const std::vector<T> &v) {
      return static_cast<int>(std::upper_bound(v.begin(), v.end(), x) - v.begin());
    });
  }

  int count_in(int lo, int hi, const T &x, const T &y) const {
    assert(0 <= lo && lo <= hi && hi < len);
    assert(!(y < x));
    return query(1, 0, len - 1, lo, hi, [&](const std::vector<T> &v) {
      return static_cast<int>(
          std::upper_bound(v.begin(), v.end(), y) - std::lower_bound(v.begin(), v.end(), x)
      );
    });
  }
};

/*** Example Usage ***/

using namespace std;

int main() {
  vector<int> a{5, 2, 8, 6, 1, 9, 3};
  MergeSortTree<int> t(a);

  assert(t.size() == 7);
  assert(t.count_leq(0, 6, 5) == 4);    // 5, 2, 1, 3
  assert(t.count_leq(2, 4, 6) == 2);    // 6, 1
  assert(t.count_in(0, 6, 3, 8) == 4);  // 5, 8, 6, 3
  assert(t.count_in(1, 5, 2, 6) == 2);  // 2, 6 (8, 1, 9 excluded)
  return 0;
}