Weighted Segment Tree in Fenwick Tree

The problem of Static Range k-th Smallest Value (POJ 2104 K-th Number) can be solved using a Weighted Segment Tree in \(O(n\log n)\) time complexity.

What if the range becomes dynamic? That is, if we also need to support an operation: modifying the value at a single position, what should we do?

Example Problem Second-Order Balanced Tree (Tree of Trees)

Maintain an ordered sequence, which requires the following operations:

  • Query the rank of \(x\) within the interval;
  • Query the value with rank \(k\) within the interval;
  • Modify the value at a certain position;
  • Query the predecessor of \(x\) within the interval (predecessor defined as the largest number less than \(x\));
  • Query the successor of \(x\) within the interval (successor defined as the smallest number greater than \(x\)).
Example Problem Luogu P2617 Dynamic Rankings

Given a sequence \(a_1,a_2 \dots a_n\) containing \(n\) numbers, support two types of operations:

  • Q l r k means query the \(k\)-th smallest number with indices in the interval \([l,r]\).
  • C x y means change \(a_x\) to \(y\).

If we use the method discussed in Segment Tree of Balanced Trees, i.e., for each node of the segment tree, maintain a balanced tree for the interval it represents, and then use binary search to find the \(k\)-th smallest value. Since each query operation covers multiple intervals, i.e., multiple nodes, but balanced trees cannot search multiple values simultaneously, the time complexity is \(O(n\log^3 n)\), which is not optimal.

The optimization idea is to combine the binary search for the answer with the operation of querying the count of numbers less than a certain value, using a Segment Tree of Dynamically Allocated Weighted Segment Trees. Since all segment trees have the same structure, we can perform segment tree binary search on multiple trees simultaneously.

During a modification operation, we first traverse down the segment tree to the point being modified, delete the old value from the dynamically allocated weighted segment trees pointed to by the nodes along the path, and then insert the new value. This involves \(O(\log n)\) segment tree nodes, and each modification in the dynamically allocated weighted segment tree is \(O(\log n)\), so the time complexity of a modification operation is \(O(\log^2 n)\).

When querying for an answer, we first collect all the segment tree nodes that cover the interval, then use a method similar to the static range \(k\)-th smallest value to move these nodes together to the left child or right child. If the total value stored in the left children of all these nodes is greater than or equal to \(k\), we move left; otherwise, we move right. Since at most \(O(\log n)\) nodes are covered, at most that many nodes move down at each step, resulting in a time complexity of \(O(\log^2 n)\).

Due to the large constant factor of segment trees, in practice, a Fenwick Tree with a smaller constant and more convenient handling of prefix sums is often used. Additionally, the space complexity is \(O(n\log^2 n)\), so pay attention to space limits when using it.

Here is one code implementation:

Implementation
  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
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
#include <algorithm>
#include <cstdio>
#include <cstring>
#include <map>
#include <set>
#define LC o << 1
#define RC o << 1 | 1
using namespace std;
constexpr int MAXN = 1000010;
int n, m, a[MAXN], u[MAXN], x[MAXN], l[MAXN], r[MAXN], k[MAXN], cur, cur1, cur2,
    q1[MAXN], q2[MAXN], v[MAXN];
char op[MAXN];
set<int> ST;
map<int, int> mp;

struct segment_tree  // Encapsulated dynamically allocated weighted segment tree
{
  int cur, rt[MAXN * 4], sum[MAXN * 60], lc[MAXN * 60], rc[MAXN * 60];

  void build(int& o) { o = ++cur; }

  void print(int o, int l, int r) {
    if (!o) return;
    if (l == r && sum[o]) printf("%d ", l);
    int mid = (l + r) >> 1;
    print(lc[o], l, mid);
    print(rc[o], mid + 1, r);
  }

  void update(int& o, int l, int r, int x, int v) {
    if (!o) o = ++cur;
    sum[o] += v;
    if (l == r) return;
    int mid = (l + r) >> 1;
    if (x <= mid)
      update(lc[o], l, mid, x, v);
    else
      update(rc[o], mid + 1, r, x, v);
  }
} st;

// Fenwick Tree implementation
namepace fenwick_impl {
  int lowbit(int o) { return (o & (-o)); }

  void upd(int o, int x, int v) {
    for (; o <= n; o += lowbit(o)) st.update(st.rt[o], 1, n, x, v);
  }

  void gtv(int o, int* A, int& p) {
    p = 0;
    for (; o; o -= lowbit(o)) A[++p] = st.rt[o];
  }

  int qry(int l, int r, int k) {
    if (l == r) return l;
    int mid = (l + r) >> 1, siz = 0;
    for (int i = 1; i <= cur1; i++) siz += st.sum[st.lc[q1[i]]];
    for (int i = 1; i <= cur2; i++) siz -= st.sum[st.lc[q2[i]]];
    // printf("j %d %d %d %d\n",cur1,cur2,siz,k);
    if (siz >= k) {
      for (int i = 1; i <= cur1; i++) q1[i] = st.lc[q1[i]];
      for (int i = 1; i <= cur2; i++) q2[i] = st.lc[q2[i]];
      return qry(l, mid, k);
    } else {
      for (int i = 1; i <= cur1; i++) q1[i] = st.rc[q1[i]];
      for (int i = 1; i <= cur2; i++) q2[i] = st.rc[q2[i]];
      return qry(mid + 1, r, k - siz);
    }
  }
}
using namespace fenwick_impl;

// Segment Tree implementation
namespace segtree_impl {
void build(int o, int l, int r) {
  st.build(st.rt[o]);
  if (l == r) return;
  int mid = (l + r) >> 1;
  build(LC, l, mid);
  build(RC, mid + 1, r);
}

void print(int o, int l, int r) {
  printf("%d %d:", l, r);
  st.print(st.rt[o], 1, n);
  printf("\n");
  if (l == r) return;
  int mid = (l + r) >> 1;
  print(LC, l, mid);
  print(RC, mid + 1, r);
}

void update(int o, int l, int r, int q, int x, int v) {
  st.update(st.rt[o], 1, n, x, v);
  if (l == r) return;
  int mid = (l + r) >> 1;
  if (q <= mid)
    update(LC, l, mid, q, x, v);
  else
    update(RC, mid + 1, r, q, x, v);
}

void getval(int o, int l, int r, int ql, int qr) {
  if (l > qr || r < ql) return;
  if (ql <= l && r <= qr) {
    q[++cur] = st.rt[o];
    return;
  }
  int mid = (l + r) >> 1;
  getval(LC, l, mid, ql, qr);
  getval(RC, mid + 1, r, ql, qr);
}

int query(int l, int r, int k) {
  if (l == r) return l;
  int mid = (l + r) >> 1, siz = 0;
  for (int i = 1; i <= cur; i++) siz += st.sum[st.lc[q[i]]];
  if (siz >= k) {
    for (int i = 1; i <= cur; i++) q[i] = st.lc[q[i]];
    return query(l, mid, k);
  } else {
    for (int i = 1; i <= cur; i++) q[i] = st.rc[q[i]];
    return query(mid + 1, r, k - siz);
  }
}
}  // namespace segtree_impl

int main() {
  scanf("%d%d", &n, &m);
  for (int i = 1; i <= n; i++) scanf("%d", a + i), ST.insert(a[i]);
  for (int i = 1; i <= m; i++) {
    scanf(" %c", op + i);
    if (op[i] == 'C')
      scanf("%d%d", u + i, x + i), ST.insert(x[i]);
    else
      scanf("%d%d%d", l + i, r + i, k + i);
  }
  for (set<int>::iterator it = ST.begin(); it != ST.end(); it++)
    mp[*it] = ++cur, v[cur] = *it;
  for (int i = 1; i <= n; i++) a[i] = mp[a[i]];
  for (int i = 1; i <= m; i++)
    if (op[i] == 'C') x[i] = mp[x[i]];
  n += m;
  // build(1,1,n);
  for (int i = 1; i <= n; i++) upd(i, a[i], 1);
  // print(1,1,n);
  for (int i = 1; i <= m; i++) {
    if (op[i] == 'C') {
      upd(u[i], a[u[i]], -1);
      upd(u[i], x[i], 1);
      a[u[i]] = x[i];
    } else {
      gtv(r[i], q1, cur1);
      gtv(l[i] - 1, q2, cur2);
      printf("%d\n", v[qry(1, n, k[i])]);
    }
  }
  return 0;
}