Introduction

If you’re into competitive programming, you’ve probably come across this fantastic blog: https://mirror.codeforces.com/blog/entry/18051. It introduces a super efficient, compact segment tree that really helped me get a better grasp on how segment trees work.

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.

A quick Recap on 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 https://mirror.codeforces.com/blog/entry/134777), 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:

#include <functional>
#include <vector>

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(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 &mdash; 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.