GCJ/2008/1B/C/segment_tree.go

package main

// SegmentTree ...
type SegmentTree []int

// NewSegmentTree ...
func NewSegmentTree(n int) SegmentTree {
	var st SegmentTree = make([]int, 2*n)
	for i := n; i < 2*n; i++ {
		st[i] = 1
	}
	for i := n - 1; 0 < i; i-- {
		st[i] = st[2*i] + st[2*i+1]
	}
	return st
}

// Update ...
func (st SegmentTree) Update(i, val int) {
	i += len(st) / 2
	st[i] = val
	for i > 1 {
		i /= 2
		st[i] = st[2*i] + st[2*i+1]
	}
}

// Search ...
func (st SegmentTree) Search(beg, val int) (end int) {
	val %= st[1]  // Deal with the loop
	if val == 0 { // If val is zero, set val to st[1]
		val = st[1]
	}
	n := len(st) / 2
	slotCnt := st.Count(beg, n)
	if slotCnt < val { // If slot cnt is not enough, start from the begining
		val -= slotCnt
		beg = 0
	}
	beg, end = beg+n, beg+n
	for st[end] < val {
		if end%2 == 1 { // If is right child, move to left child
			val -= st[end]
			end++
		}
		end /= 2
	} // Find the interval that contains end from bottom to top
	for end < n { // While end is not leaf
		lchild, rchild := end*2, end*2+1
		if st[lchild] < val {
			val -= st[lchild]
			end = rchild // Find in right tree
		} else { // st[lchild] == val
			end = lchild // Find in left tree
		}
	}
	return end - n
}

// Count ...
func (st SegmentTree) Count(left, right int) (sum int) {
	n := len(st) / 2
	left, right = left+n, right+n
	for left < right {
		if left%2 == 1 {
			sum += st[left]
			left++
		}
		if right%2 == 1 { // right is not included, so --
			right--
			sum += st[right]
		}
		left, right = left/2, right/2
	}
	return
}