If you’re into competitive programming, you’ve probably come across this fantastic blog. It introduces a super efficient, compact segment tree really shifted the way I thought about segment trees.

So today, while experimenting with this idea, I thought: "Why not extend it to support binary search?" And that’s exactly what I did. Here’s my attempt at it, and I thought it was worth sharing.

Quick recap: std::partition_point

If you’re familiar with binary search, you might know about std::partition_point. It’s a C++ function that returns the first false value in a range. C++20 also added std::ranges::partition_point (see one-line binary search), which makes binary search even more concise: basically, a one-liner.

Taking some inspiration from this (mainly for the name), here's how I adapted it for the segment tree:

Implementation

template <typename T> class SegmentTree {
public:
  int n;
  T idt;
  std::vector<T> seg;
  std::function<T(T, T)> f;

  SegmentTree(int _n, std::function<T(T, T)> f = std::plus<T>(), T idt = T())
    : n(std::bit_ceil(uint32_t(_n))), idt(idt), f(f), seg(2 * n, idt) {}

  void set(int idx, T x) {
    for (seg[idx += n] = x, idx /= 2; idx > 0; idx /= 2) {
      seg[idx] = f(seg[2 * idx], seg[2 * idx + 1]);
    }
  }

  T query(int l, int r) {
    T ansL = idt, ansR = idt;
    for (l += n, r += n + 1; l < r; l /= 2, r /= 2) {
      if (l & 1) ansL = f(ansL, seg[l++]);
      if (r & 1) ansR = f(seg[--r], ansR);
    }
    return f(ansL, ansR);
  }

  int partition_point(int l, const std::function<bool(T)> &t) {
    T p = idt;
    for (l += n; t(f(p, seg[l])) and l; l /= 2) {
      if (l & 1 and l != 1) p = f(p, seg[l++]);
    }
    if (!l) {
      return n;
    }
    while (l < n) {
      if (t(f(p, seg[l <<= 1]))) p = f(p, seg[l++]);
    }
    return l - n;
  }
};

The partition_point() function

This is where I added binary search. The partition_point() function finds the first index in the range $$$[l, n)$$$ where the predicate function t returns false by simulating a traversal of the segment tree.

Why share this?

I know there are other segment tree implementations out there, like AtCoder’s library, but I think this one is a bit cleaner and looks more elegant.

It’s not groundbreaking, but I thought I’d share it anyway. Maybe someone out there will find it useful.