BIT.cpp

// Binary indexed tree supporting binary search.
struct BIT {
  int n;
  vector<int> bit;
  // BIT can be thought of as having entries f[1], ..., f[n]
  // which are 0-initialized
  BIT(int n) : n(n), bit(n + 1) {}
  // returns f[1] + ... + f[idx-1]
  // precondition idx <= n+1
  int read(int idx) {
    idx--;
    int res = 0;
    while (idx > 0) {
      res += bit[idx];
      idx -= idx & -idx;
    }
    return res;
  }
  // returns f[idx1] + ... + f[idx2-1]
  // precondition idx1 <= idx2 <= n+1
  int read2(int idx1, int idx2) {
    return read(idx2) - read(idx1);
  }
  // adds val to f[idx]
  // precondition 1 <= idx <= n (there is no element 0!)
  void update(int idx, int val) {
    while (idx <= n) {
      bit[idx] += val;
      idx += idx & -idx;
    }
  }
  // returns smallest positive idx such that read(idx) >=
  // target
  int lower_bound(int target) {
    if (target <= 0)
      return 1;
    int pwr = 1;
    while (2 * pwr <= n)
      pwr *= 2;
    int idx = 0;
    int tot = 0;
    for (; pwr; pwr >>= 1) {
      if (idx + pwr > n)
        continue;
      if (tot + bit[idx + pwr] < target) {
        tot += bit[idx += pwr];
      }
    }
    return idx + 2;
  }
  // returns smallest positive idx such that read(idx) >
  // target
  int upper_bound(int target) {
    if (target < 0)
      return 1;
    int pwr = 1;
    while (2 * pwr <= n)
      pwr *= 2;
    int idx = 0;
    int tot = 0;
    for (; pwr; pwr >>= 1) {
      if (idx + pwr > n)
        continue;
      if (tot + bit[idx + pwr] <= target) {
        tot += bit[idx += pwr];
      }
    }
    return idx + 2;
  }
};