Artifact b7ca5af5f1eab32acf65d6df108247dc176ede43

• File lib/dataStructure/segmentTree.cpp
  • 2013-09-14 20:29:31 - part of checkin [94afd8e13c] on branch trunk - Segment tree with range update. (user: kinaba) [annotate]

//-------------------------------------------------------------
// Segment tree
//   in some general form
//
// Verified by
//   - Codeforces 200 D
//   - Codeforces 107 C (old version)
//   - Codeforces 104 E (old version)
//-------------------------------------------------------------

class SegmentTree
{
	struct Node
	{
		int sum;

		static Node Zero()
		{
			Node c = {0};
			return c;
		}
		static Node One(int v)
		{
			Node c = {v};
			return c;
		}
		static Node Concat(const Node& l, const Node& r)
		{
			Node c = {l.sum + r.sum};
			return c;
		}
		static Node Repeat(const Node& n, int k)
		{
			if(k==0) return Zero();
			Node c = {n.sum * k};
			return c;
		}

		bool lazy;
	};

public:
	template<typename Seq>
	SegmentTree(const Seq& s) {
		int N = 1;
		while( N < s.size() )
			N <<= 1;

		tree.resize(tree.size()+1); tree.back().resize(N);
		for(int i=0; i<N; ++i)
			tree.back()[i] = i<s.size() ? Node::One(s[i]) : Node::Zero();

		while(N>>=1) {
			tree.resize(tree.size()+1); tree.back().resize(N);
			for(int i=0; i<N; ++i)
				CalcMidNode(tree.size()-1, i);
		}
	}

	Node Query(int s, int e) { // compute h( seq[s,e) ) :  O(log n)
		return QueryRec(s, e, tree.size()-1, 0, tree[0].size());
	}

	template<typename Value>
	void Set(int s, int e, Value v) { // seq[s,e):=v : O(log n)
		SetRec(s, e, Node::One(v), tree.size()-1, 0, tree[0].size());
	}

private:
	Node QueryRec(int s, int e, int lv, int idx, int stride)
	{
		const int myL = stride*idx;
		const int myR = stride*(idx+1);
		if( e<=myL || myR<=s )
			return Node::Zero();
		ResolveLazy(lv, idx);

		if( s<=myL && myR<=e )
			return tree[lv][idx];
		return Node::Concat(QueryRec(s,e,lv-1,idx*2,stride/2),
		                    QueryRec(s,e,lv-1,idx*2+1,stride/2));
	}

	void SetRec(int s, int e, const Node& n, int lv, int idx, int stride)
	{
		const int myL = stride*idx;
		const int myR = stride*(idx+1);
		if( e<=myL || myR<=s )
			return;
		ResolveLazy(lv, idx);

		if( stride == 1 ) {
			tree[lv][idx] = n;
		} else {
			if( s<=myL && myR<=e ) {
				tree[lv][idx] = n;
				tree[lv][idx].lazy = true;
				ResolveLazy(lv, idx);
				return;
			}
			SetRec(s,e,n,lv-1,idx*2,stride/2);
			SetRec(s,e,n,lv-1,idx*2+1,stride/2);
			CalcMidNode(lv, idx);
		}
	}

	void CalcMidNode(int lv, int idx)
	{
		tree[lv][idx] = Node::Concat(tree[lv-1][idx*2], tree[lv-1][idx*2+1]);
	}

	void ResolveLazy(int lv, int idx)
	{
		if(tree[lv][idx].lazy) {
			if(lv > 0) {
				tree[lv-1][idx*2]   = tree[lv][idx];
				tree[lv-1][idx*2+1] = tree[lv][idx];
			}
			tree[lv][idx] = Node::Repeat(tree[lv][idx], 1<<lv);
		}
	}

	vector< vector<Node> > tree;
};