disjoint_set/segment_tree.cpp

#include <iostream>
#include "segment_tree.hpp"

#define endl '\n'

using namespace std;

void example1()
{
    cout << "\nExample 1" << endl;
    // Segment tree that allows setting an element in a range and fetching a single element.
    auto st = segment_tree<int, int>()
                  ->with_merge([&](int &a, int &b) {
                      return 0;
                  })
                  ->with_lazy([&](int &a, int &b, int &c) {
                      if (a)
                      {
                          b = a;
                          c = a;
                      }
                  })
                  ->with_update([&](int &a, int &b) {
                      a = b;
                  })
                  ->init(4);

    st->update(0, 1);
    cout << "Set value 1 at position 0" << endl;

    st->update(1, 2);
    cout << "Set value 2 at position 1" << endl;

    cout << "All values" << endl;
    for (int i = 0; i < 4; ++i)
        cout << st->query(i) << " ";
    cout << endl;

    cout << "Set value 4 in the range [1,3)" << endl;
    st->update(1, 3, 4);

    cout << "All values" << endl;
    for (int i = 0; i < 4; ++i)
        cout << st->query(i) << " ";
    cout << endl;
}

void example2()
{
    cout << "\nExample 2" << endl;
    // Segment tree of maximum that allows setting adding in a point
    // It stores a pair on each node:
    //      - Maximum on the range
    //      - Lazy to set children
    auto st = segment_tree<int, int>()
                  ->with_merge([&](int &a, int &b) {
                      return max(a, b);
                  })
                  ->with_update([&](int &a, int &b) {
                      a += b;
                  })
                  ->init(4);

    st->update(0, 1);

    st->update(1, 5);

    for (int i = 0; i < 4; ++i)
        cout << st->query(i) << " ";
    cout << endl;

    st->update(1, 3);
    st->update(2, 2);

    for (int i = 1; i <= 4; ++i)
    {
        for (int j = 0; j + i <= 4; ++j)
            cout << st->query(j, j + i) << " ";
        cout << endl;
    }
}

void example3()
{
    cout << "\nExample 3" << endl;
    // Segment tree of maximum that allows setting a value on range
    // It stores a pair on each node:
    //      - Maximum on the range
    //      - Lazy to set children
    typedef pair<int, int> T;
    auto st = segment_tree<T, int>()
                  ->with_merge([&](T &a, T &b) {
                      return make_pair(max(a.first, b.first), 0);
                  })
                  ->with_lazy([&](T &a, T &b, T &c) {
                      if (a.second)
                      {
                          b = make_pair(a.second, a.second);
                          c = make_pair(a.second, a.second);
                      }
                  })
                  ->with_update([&](T &a, int &b) {
                      a = make_pair(b, b);
                  })
                  ->init(4);

    st->update(0, 1);

    st->update(1, 5);

    for (int i = 0; i < 4; ++i)
        cout << st->query(i).first << " ";
    cout << endl;

    st->update(1, 4, 3);
    st->update(2, 4, 2);

    for (int i = 1; i <= 4; ++i)
    {
        for (int j = 0; j + i <= 4; ++j)
            cout << st->query(j, j + i).first << " ";
        cout << endl;
    }
}
int main()
{
    example1();
    example2();
    example3();

    return 0;
}