How can I implement segment trees with lazy propagation?

Question

I am implementing a segment tree, to be able to answer quickly to the following queries in an array A:

• query i, j: sum of all elements in range (i,j)
• update i, j, k: add k to all elements in the range(i,j)

Here is my implementation:

typedef long long intt;

const int max_num=100000,max_tree=4*max_num;
intt A[max_num],ST[max_tree];

void initialize(int node, int be, int en) {
  if(be==en) {
    ST[node]=ST[be];
  } else {
    initialize(2*node+1,be,(be+en)/2);
    initialize(2*node+2,(be+en)/2+1,en);

    ST[node]=ST[2*node+1]+ST[2*node+2];
  }
}

void upg(int node, int be, int en, int i, intt k) {
  if(be>i || en<i || be>en) return;
  if(be==en) {
    ST[node]+=k;
    return;
  }
  upg(2*node+1, be, (be+en)/2, i, k);
  upg(2*node+2, (be+en)/2+1, en, i, k);
  ST[node] = ST[2*node+1]+ST[2*node+2];
}

intt query(int node, int be, int en, int i, int j) {
  if(be>j || en<i) return -1;
  if(be>=i && en<=j) return ST[node];

  intt q1=query(2*node+1, be, (be+en)/2, i, j);
  intt q2=query(2*node+2, (be+en)/2+1, en, i, j);

  if(q1==-1) return q2;
  else if(q2==-1) return q1;
  else return q1+q2;
}

The query function is really fast, its complexity is O(lg N), where N is j-i. The update function is also fast in the average case, but when j-i is big, the complexity of the update is O(N lg N), which is not fast at all.

I have searched the subject a bit, and I found that if I implement the segment tree with lazy propagation, the complexity of both query and update is O(lg N), which is asymptotically faster than O(N lg N).

I also found a link to an other question, which has a really nice implementation of a segment tree, that uses pointers: How to implement segment trees with lazy propagation?. So, here is my question: Is there a simpler way to implement lazy propagation, without using pointers, but with array indexes, and without a segment_tree data structure?

Answer 1

This is me playing around with this data structure and some template tomfoolery.

At the bottom of all of this mess is accessing two flat arrays, one of them containing a tree of sums, the other containing a tree of carry values to propogate down later. Conceptually they form one binary tree.

The true value of a node in the binary tree is the value in stored sum tree, plus the number of leaves under the node times the sum of all carry tree values from the node back up to the root.

At the same time, the true value of each node in the tree is equal to the true value of each leaf node under it.

I wrote one function to do both carry and sum, because it turns out they are visiting the same nodes. And a read sometimes writes. So you get a sum by calling it with an increase of zero.

All the template tomfoolery does is do the math for where the offsets into each tree the nodes are, and where the left and right children are.

While I do use a struct, the struct is transient -- it is just a wrapper with some pre calculated values around an offset into an array. I do store a pointer to the start of the array, but every block_ptr uses the exact same root value in this program.

For debugging, I have some craptacular Assert() and Debug() macros, plus a trace nullary function that the recursive sum function calls (which I use to track the total number of calls into it). Once again, being needlessly complex to avoid global state. :)

#include <memory>
#include <iostream>

// note that you need more than 2^30 space to fit this
enum {max_tier = 30};

typedef long long intt;

#define Assert(x) (!(x)?(std::cout << "ASSERT FAILED: (" << #x << ")\n"):(void*)0)
#define DEBUG(x) 

template<size_t tier, size_t count=0>
struct block_ptr
{
  enum {array_size = 1+block_ptr<tier-1>::array_size * 2};
  enum {range_size = block_ptr<tier-1>::range_size * 2};
  intt* root;
  size_t offset;
  size_t logical_offset;
  explicit block_ptr( intt* start, size_t index, size_t logical_loc=0 ):root(start),offset(index), logical_offset(logical_loc) {}
  intt& operator()() const
  {
    return root[offset];
  }
  block_ptr<tier-1> left() const
  {
    return block_ptr<tier-1>(root, offset+1, logical_offset);
  }
  block_ptr<tier-1> right() const
  {
    return block_ptr<tier-1>(root, offset+1+block_ptr<tier-1>::array_size, logical_offset+block_ptr<tier-1>::range_size);
  }
  enum {is_leaf=false};
};

template<>
struct block_ptr<0>
{
  enum {array_size = 1};
  enum {range_size = 1};
  enum {is_leaf=true};
  intt* root;
  size_t offset;
  size_t logical_offset;

  explicit block_ptr( intt* start, size_t index, size_t logical_loc=0 ):root(start),offset(index), logical_offset(logical_loc)
  {}
  intt& operator()() const
  {
    return root[offset];
  }
  // exists only to make some of the below code easier:
  block_ptr<0> left() const { Assert(false); return *this; }
  block_ptr<0> right() const { Assert(false); return *this; }
};


template<size_t tier>
void propogate_carry( block_ptr<tier> values, block_ptr<tier> carry )
{
  if (carry() != 0)
  {
    values() += carry() * block_ptr<tier>::range_size;
    if (!block_ptr<tier>::is_leaf)
    {
      carry.left()() += carry();
      carry.right()() += carry();
    }
    carry() = 0;
  }
}

// sums the values from begin to end, but not including end!
// ie, the half-open interval [begin, end) in the tree
// if increase is non-zero, increases those values by that much
// before returning it
template<size_t tier, typename trace>
intt query_or_modify( block_ptr<tier> values, block_ptr<tier> carry, int begin, int end, int increase=0, trace const& tr = [](){} )
{
  tr();
  DEBUG(
  std::cout << begin << " " << end << " " << increase << "\n";
  if (increase)
  {
    std::cout << "Increasing " << end-begin << " elements by " << increase << " starting at " << begin+values.offset << "\n";
  }
  else
  {
    std::cout << "Totaling " << end-begin << " elements starting at " << begin+values.logical_offset << "\n";
  }
  )
  if (end <= begin)
    return 0;
  size_t mid = block_ptr<tier>::range_size / 2;
  DEBUG( std::cout << "[" << values.logical_offset << ";" << values.logical_offset+mid << ";" << values.logical_offset+block_ptr<tier>::range_size << "]\n"; )
  // exatch math first:
  bool bExact = (begin == 0 && end >= block_ptr<tier>::range_size);
  if (block_ptr<tier>::is_leaf)
  {
    Assert(bExact);
  }
  bExact = bExact || block_ptr<tier>::is_leaf; // leaves are always exact
  if (bExact)
  {
    carry()+=increase;
    intt retval =  (values()+carry()*block_ptr<tier>::range_size);
    DEBUG( std::cout << "Exact sum is " << retval << "\n"; )
    return retval;
  }
  // we don't have an exact match.  Apply the carry and pass it down to children:
  propogate_carry(values, carry);
  values() += increase * end-begin;
  // Now delegate to children:
  if (begin >= mid)
  {
    DEBUG( std::cout << "Right:"; )
    intt retval = query_or_modify( values.right(), carry.right(), begin-mid, end-mid, increase, tr );
    DEBUG( std::cout << "Right sum is " << retval << "\n"; )
    return retval;
  }
  else if (end <= mid)
  {
    DEBUG( std::cout << "Left:"; )
    intt retval = query_or_modify( values.left(), carry.left(), begin, end, increase, tr );
    DEBUG( std::cout << "Left sum is " << retval << "\n"; )
    return retval;
  }
  else
  {
    DEBUG( std::cout << "Left:"; )
    intt left = query_or_modify( values.left(), carry.left(), begin, mid, increase, tr );
    DEBUG( std::cout << "Right:"; )
    intt right = query_or_modify( values.right(), carry.right(), 0, end-mid, increase, tr );
    DEBUG( std::cout << "Right sum is " << left << " and left sum is " << right << "\n"; )
    return left+right;
  }
}

Here are some helper classes to make creating a segment tree of a given size easy. Note, however, that all you need is an array of the right size, and you can construct a block_ptr from a pointer to element 0, and you are good to go.

template<size_t tier>
struct segment_tree
{
  typedef block_ptr<tier> full_block_ptr;
  intt block[full_block_ptr::range_size];
  full_block_ptr root() { return full_block_ptr(&block[0],0); }
  void init()
  {
    std::fill_n( &block[0], size_t(full_block_ptr::range_size), 0 );
  }
};

template<size_t entries, size_t starting=0>
struct required_tier
{
  enum{ tier =
    block_ptr<starting>::array_size >= entries
    ?starting
    :required_tier<entries, starting+1>::tier
  };
  enum{ error =
    block_ptr<starting>::array_size >= entries
    ?false
    :required_tier<entries, starting+1>::error
  };
};

// max 2^30, to limit template generation.
template<size_t entries>
struct required_tier<entries, size_t(max_tier)>
{
  enum{ tier = 0 };
  enum{ error = true };
};

// really, these just exist to create an array of the correct size
typedef required_tier< 1000000 > how_big;

enum {tier = how_big::tier};


int main()
{
  segment_tree<tier> values;
  segment_tree<tier> increments;
  Assert(!how_big::error); // can be a static assert -- fails if the enum of max tier is too small for the number of entries you want
  values.init();
  increments.init();
  auto value_root = values.root();
  auto carry_root = increments.root();

  size_t count = 0;
  auto tracer = [&count](){count++;};
  intt zero = query_or_modify( value_root, carry_root, 0, 100000, 0, tracer );
  std::cout << "zero is " << zero << " in " << count << " steps\n";
  count = 0;
  Assert( zero == 0 );
  intt test2 = query_or_modify( value_root, carry_root, 0, 100, 10, tracer ); // increase everything from 0 to 100 by 10
  Assert(test2 == 1000);
  std::cout << "test2 is " << test2 << " in " << count << " steps \n";
  count = 0;
  intt test3 = query_or_modify( value_root, carry_root, 1, 1000, 0, tracer );
  Assert(test3 == 990);
  std::cout << "test3 is " << test3 << " in " << count << " steps\n";
  count = 0;
  intt test4 = query_or_modify( value_root, carry_root, 50, 5000, 87, tracer );
  Assert(test4 == 10*(100-50) + 87*(5000-50) );
  std::cout << "test4 is " << test4 << " in " << count << " steps\n";
  count = 0;
}

While this isn't the answer you want, it might make it easier for someone to write it. And writing this amused me. So, hope it helps!

The code was tested and compiled on Ideone.com using a C++0x compiler.