avighnakc's blog

Efficient iterative segment tree with binary search
By avighnakc, history, 17 months ago, In English

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, 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), 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 + 1); l /= 2) {
      if (l & 1) p = f(p, seg[l++]);
    }
    if (t(f(p, seg[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.

  • Vote: I like it
  • +38
  • Vote: I do not like it

»
17 months ago, hide # |
← Rev. 3  
Vote: I like it +3 Vote: I do not like it

Additional information: here's an example AC submission to CSES Hotel Queries.

(Note: ignore the first and second revision of this comment. That simplification turned out to be incorrect)

→ Reply
»
17 months ago, hide # |
 
Vote: I like it -20 Vote: I do not like it

This already exists and is well known: https://usaco.guide/plat/segtree-ext#walking-on-a-segment-tree

→ Reply
»
17 months ago, hide # |
 
Vote: I like it -13 Vote: I do not like it

BTW, This is known as walking on a segment tree

→ Reply
»
17 months ago, hide # |
 
Vote: I like it 0 Vote: I do not like it

https://github.com/atcoder/ac-library/blob/master/atcoder/segtree.hpp

→ Reply
»
17 months ago, hide # |
 
Vote: I like it 0 Vote: I do not like it

Personally, it is worth making all segment trees perfect power of 2. Descent code (your ST-walking) is simpler than binary search and custom segtree operations are easier to write.

→ Reply
»
16 months ago, hide # |
← Rev. 2  
Vote: I like it 0 Vote: I do not like it

You can still walk without setting $$$n$$$ to the nearest $$$2^k$$$ as shown here (for those who are curious).

P.S. In your blog the early return condition is (l & -l) == l instead of the correct !l.

→ Reply
  • »
    »
    16 months ago, hide # ^ |
    ← Rev. 2  
    Vote: I like it 0 Vote: I do not like it

    I'm pretty sure the latter breaks when you use partition point with anything other than 0.

    The same is true for the bit_ceil thing. It is necessary for correct functionality when you do any query that isn't at 0.

    → Reply
    • »
      »
      »
      16 months ago, hide # ^ |
      ← Rev. 2  
      Vote: I like it 0 Vote: I do not like it

      Oops I just realized I made a mistake in copy your code over to the editor (I don't like copy & pasting) and made the condition if (l & -l) instead. Which made me confused on why it did not work. My bad.

      Thanks for clarifying though!

      → Reply
    • »
      »
      »
      16 months ago, hide # ^ |
      ← Rev. 2  
      Vote: I like it 0 Vote: I do not like it

      Seems like it's not working for $$$l = 0$$$

      8
1 2 3 4 5 6 7 8

      Expected $$$5$$$, output $$$8$$$. And if you change the condition to !l you get $$$2$$$.

      code
      → Reply
      • »
        »
        »
        »
        16 months ago, hide # ^ |
        ← Rev. 3  
        Vote: I like it 0 Vote: I do not like it

        Okay, so $$$2$$$ is correct. I guess you meant to make it a max segment tree but left it as a sum one. You are right that $$$8$$$ is wrong though. Thanks for the counter-example.

        Current state is: if you use !l, you'll always get the right answer with l = 0, and you get the right answer if the answer is not $$$n$$$ (otherwise you get $$$0$$$, which I'm trying to fix right now).

        Edit: I seemed to have found why that was happening. Counter-example if you're curious:

        int main() {
  int n = 10;
  std::vector<int> a = {20, 15, 17, 35, 25, 40, 12, 19, 13, 12};

  const int inf = 1e9;

  SegmentTree<int> st(n, [](int a, int b) { return std::max(a, b); }, -inf);
  for (int i = 0, x; i < n; i++) {
    st.set(i, a[i]);
  }
  cout << st.partition_point(8, [](int mx) { return mx < 19; }) << "\n";
}

        Expected $$$10$$$, actual $$$0$$$. Turns out the mistake was that in my for loop, I had to check if I was at the "right boundary" of the segment tree, in which case l++ was invalid. To check that I just had to switch to checking whether $$$l + 1$$$ is a power of $$$2$$$ in the loop. And the correct condition for returning $$$n$$$ also becomes:

        if (t(f(p, seg[l]))) {
  return _n;
}

        I've corrected my implementation in the post. Here's an AC submission on CSES Mountain Range which really does test it properly on values $$$l$$$ values that are not just $$$0$$$.

        PS: Thanks a lot for helping me fix it!

        → Reply