Binary Tree to set/add/sum/max

binary.cpp

#include <iostream>
#include <string>
#include <algorithm>
#include <limits>
using namespace std;

// The height of the tree (ok, the name is counter-intuitive...)
#define DEPTH 20

// The 'undefined' value for the 'set' field
#define UNDEF numeric_limits<long long>::min()

struct node
{
    long long set;
    long long add;
    long long sum;
    long long max;
};

node tree[1 << (DEPTH + 1)];

int left_child (int index)
{
    return 2 * index;
}

int right_child (int index)
{
    return 2 * index + 1;
}

int get_left (int index)
{
    int top = sizeof(int) * 8 - __builtin_clz(index) - 1;
    return (1 << (DEPTH - top)) * (index ^ (1 << top));
}

int get_right (int index)
{
    int top = sizeof(int) * 8 - __builtin_clz(index) - 1;
    return (1 << (DEPTH - top)) * ((index ^ (1 << top)) + 1);
}

// Do the actual 'set' command on a node
void set_on_node(int index, long long value)
{
    tree[index].set = value;
    tree[index].add = 0;
    tree[index].sum = value * (get_right(index) - get_left(index));
    tree[index].max = value;
}

// Do the actual 'add' command on a node
void add_on_node(int index, long long value)
{
    tree[index].add += value;
    tree[index].sum += value * (get_right(index) - get_left(index));
    tree[index].max += value;
}

void lazy_update (int index)
{
    if (tree[index].set != UNDEF)
    {
        set_on_node(left_child(index), tree[index].set);
        set_on_node(right_child(index), tree[index].set);
        tree[index].set = UNDEF;
    }
    add_on_node(left_child(index), tree[index].add);
    add_on_node(right_child(index), tree[index].add);
    tree[index].add = 0;
}

// Set [a, b) to value
void set (int index, int a, int b, long long value)
{
    int l = get_left(index);
    int r = get_right(index);
    // [a, b) contains the node
    if (a <= l && r <= b)
    {
        set_on_node(index, value);
    }
    // [a, b) doesn't intersect the node
    else if (b <= l || r <= a)
    {
        // do nothing
    }
    else
    {
        lazy_update(index);
        set(left_child(index), a, b, value);
        set(right_child(index), a, b, value);
        tree[index].sum = tree[left_child(index)].sum + tree[right_child(index)].sum;
        tree[index].max = max(tree[left_child(index)].max, tree[right_child(index)].max);
    }
}

// Add value to [a, b)
void add (int index, int a, int b, long long value)
{
    int l = get_left(index);
    int r = get_right(index);
    // [a, b) contains the node
    if (a <= l && r <= b)
    {
        add_on_node(index, value);
    }
    // [a, b) doesn't intersect the node
    else if (b <= l || r <= a)
    {
        // do nothing
    }
    else
    {
        lazy_update(index);
        add(left_child(index), a, b, value);
        add(right_child(index), a, b, value);
        tree[index].sum = tree[left_child(index)].sum + tree[right_child(index)].sum;
        tree[index].max = max(tree[left_child(index)].max, tree[right_child(index)].max);
    }
}

// Get the sum of [a, b)
long long sum (int index, int a, int b)
{
    int l = get_left(index);
    int r = get_right(index);
    // [a, b) contains the node
    if (a <= l && r <= b)
    {
        return tree[index].sum;
    }
    // [a, b) doesn't intersect the node
    else if (b <= l || r <= a)
    {
        return 0;
    }
    else
    {
        lazy_update(index);
        return sum(left_child(index), a, b)
             + sum(right_child(index), a, b);
    }
}

// Get the max of [a, b)
// (the underscore is needed to avoid conflicts with max() from <algorithm>)
long long max_ (int index, int a, int b)
{
    int l = get_left(index);
    int r = get_right(index);
    // [a, b) contains the node
    if (a <= l && r <= b)
    {
        return tree[index].max;
    }
    // [a, b) doesn't intersect the node
    else if (b <= l || r <= a)
    {
        return UNDEF;
    }
    else
    {
        lazy_update(index);
        return max(max_(left_child(index), a, b),
                   max_(right_child(index), a, b));
    }
}

int n, m;

int main ()
{
    cin >> n >> m;

    // Read initial values
    for (int i = 0; i < n; i += 1)
    {
        cin >> tree[(1 << DEPTH) + i].set;
        tree[(1 << DEPTH) + i].add = 0;
        tree[(1 << DEPTH) + i].sum = tree[(1 << DEPTH) + i].set;
        tree[(1 << DEPTH) + i].max = tree[(1 << DEPTH) + i].set;
    }

    // Set the other nodes to their default value
    for (int i = n; i < (1 << DEPTH); i += 1)
    {
        tree[(1 << DEPTH) + i].set = UNDEF;
        tree[(1 << DEPTH) + i].add = 0;
        tree[(1 << DEPTH) + i].sum = 0;
        tree[(1 << DEPTH) + i].max = UNDEF;
    }
    for (int i = (1 << DEPTH) - 1; i > 0; i -= 1)
    {
        tree[i].set = UNDEF;
        tree[i].add = 0;
        tree[i].sum = tree[left_child(i)].sum + tree[right_child(i)].sum;
        tree[i].max = max(tree[left_child(i)].max, tree[right_child(i)].max);
    }

    string cmd;
    int a, b;
    long long value;

    for (int i = 0; i < m; i += 1)
    {
        cin >> cmd;

        if (cmd == "set")
        {
            cin >> a >> b >> value;
            set(1, a, b, value);
        }
        else if (cmd == "add")
        {
            cin >> a >> b >> value;
            add(1, a, b, value);
        }
        else if (cmd == "sum")
        {
            cin >> a >> b;
            cout << sum(1, a, b) << endl;
        }
        else // cmd == "max"
        {
            cin >> a >> b;
            cout << max_(1, a, b) << endl;
        }
    }
}