线段树-Segment Tree

问题：给定一系列区间I，一个点q，找出所有包含q的区间

问题来源：http://www.cs.nthu.edu.tw/~wkhon/ds/ds10/tutorial/tutorial6.pdf

线段树

https://en.wikipedia.org/wiki/Segment_tree

假设有n个区间，所有的区间的端点（包括左端点和右端点）至多有2n个。将这些端点排序，构成基本区间(elementary interval)：
(-∞,p1), [p1,p1],(p1,p2),[p2,p2],(p2,p3),...,(pn,+∞)，至多有4n+1个基本区间。

使用堆的存储结构即可构造一棵完全二叉树，每个节点的区间是其子节点的并集。

构造

对于每一个区间X, 将其插入节点N，

• 如果N的区间是X子集，则将X加入到N的列表中，结束；
• 否则，如果X与N的左节点或右节点有交集，则递归将X加入到左右子节点中。（注意：这里如果左右节点都有交集，X都会被加入到左右子节点中）

查询

对于点q, 使用二分查找即可确定所有包含q的集合。

代码

复杂点：1.区间的表示，区间的包含，相交判断

// valueRange 区间
type valueRange struct {
	lower    T
	upper    T
	lowerInc bool // lower inclusive
	upperInc bool // upper inclusive
}

// a,b must not be nil
func contains(a, b *valueRange) bool {
	if a.lower != nil {
		c := compareVal(b.lower, a.lower)
		if c < 0 {
			return false // b's left bound exceeds a's left bound
		}
		if c == 0 && !a.lowerInc && b.lowerInc {
			return false // lower not inclusive
		}
	}
	if a.upper != nil {
		c := compareVal(a.upper, b.upper)
		if c < 0 {
			return false // b's right bound exceeds a's right bound
		}
		if c == 0 && !a.upperInc && b.upperInc {
			return false // upper not inclusive
		}
	}
	return true
}

func containsVal(a *valueRange, v T) bool {
	// v must not be nil
	if v == nil {
		panic("containsVal v should not be nil")
	}
	// compare with lower
	d := compareVal(a.lower, v)
	if d > 0 {
		return false
	}
	if d == 0 {
		return a.lowerInc
	}
	// compare with upper
	d = compareVal(v, a.upper)
	if d > 0 {
		return false
	}
	if d == 0 {
		return a.upperInc
	}
	return true
}

// constraint: a,b at least one range is open
func intersects(a, b *valueRange) bool {
	// false condition:
	//    a.upper < b.lower or a.lower > b.upper
	if a.upper != nil {
		c := compareVal(b.lower, a.upper)
		if c == 0 {
			if a.upperInc && b.lowerInc {
				return true
			}
			return false
		}
		if c > 0 {
			return false
		}
	}
	if a.lower != nil {
		c := compareVal(a.lower, b.upper)
		if c == 0 {
			if a.lowerInc && b.upperInc {
				return true
			}
			return false
		}
		if c > 0 {
			return false
		}
	}
	return true
}

// segment tree
type segmentTree struct {
	nodes []*segmentNode
}
type segmentNode struct {
	*valueRange
	idx Bitmap
}

func buildSegmentTree(rgs []*idxRange) *segmentTree {
	// sort all points
	endpoints := make([]T, 0, len(rgs)*2)
	for _, rg := range rgs {
		if rg.lower != nil {
			endpoints = append(endpoints, rg.lower)
		}
		if rg.upper != nil {
			endpoints = append(endpoints, rg.upper)
		}
	}
	sort.Slice(endpoints, func(i, j int) bool {
		return compareVal(endpoints[i], endpoints[j]) < 0
	})
	if len(endpoints) == 0 {
		return &segmentTree{}
	}
	// remove duplicate
	j := 0
	for i := 1; i < len(endpoints); i++ {
		if compareVal(endpoints[i], endpoints[i-1]) != 0 {
			j++
			endpoints[j] = endpoints[i]
		}
	}
	endpoints = endpoints[:j+1]

	// let a be the tree from root to leaf, i denotes the depth, a[i] denotes the index in array `nodes`
	// a[i+1] = 2*a[i]+1
	// a[i+1]+1 = 2*(a[i]+1) = 2^(i+1-k)*(a[k]+1) = 2^(i+1)*(a[0]+1) = 2^(i+1)
	// a[i] = 2^i-1
	// it's height h
	//  full nodes:  2^0, 2^1, 2^2, ....
	//               s[i] = 2^0 + 2^1 + ... + 2^i
	//               2s[i] =      2^1+ ....  + 2^i  + 2^(i+1)
	//               s[i]=2s[i]-s[i] = 2^(i+1)-1
	//  x points,
	elmIntvN := 2*len(endpoints) + 1 // number of elementary intervals
	// elementary intervals
	elemIntvs := make([]*segmentNode, elmIntvN)
	leafN := upperBoundPowerOfTwo(uint64(elmIntvN))
	treeHeight := bits.Len64(leafN)
	leafIdxBegin := 1<<(treeHeight-1) - 1
	leafIdxEnd := leafIdxBegin + elmIntvN

	for i, endpoint := range endpoints {
		var lower T
		if i > 0 {
			lower = endpoints[i-1]
		}
		elemIntvs[2*i] = &segmentNode{valueRange: &valueRange{lower: lower, upper: endpoint, lowerInc: false, upperInc: false}}
		elemIntvs[2*i+1] = &segmentNode{valueRange: &valueRange{lower: endpoint, upper: endpoint, lowerInc: true, upperInc: true}}
		if i == len(endpoints)-1 {
			elemIntvs[2*i+2] = &segmentNode{valueRange: &valueRange{lower: endpoint, upper: nil, lowerInc: false, upperInc: false}}
		}
	}

	// tree example:
	//            0
	//     1             2
	//  3     4        5      6
	// 7 8   9 10    11 12  13  14
	nodes := make([]*segmentNode, leafIdxEnd)
	// traverse each node to fill
	// d: depth(starting from 1)
	// i: index in the array, i.e.  nodes[i] denotes current node
	var buildTree func(i int)
	buildTree = func(i int) {
		if i >= leafIdxEnd {
			return
		}
		// leaf node
		if i >= leafIdxBegin {
			nodes[i] = elemIntvs[i-leafIdxBegin] // elementary interval index
			return
		}
		buildTree(2*i + 1) // left child
		buildTree(2*i + 2) // right child
		var lower T
		var lowerInc bool
		var upper T
		var upperInc bool
		if 2*i+1 < leafIdxEnd {
			lower = nodes[2*i+1].lower
			lowerInc = nodes[2*i+1].lowerInc
		}
		if 2*i+2 < leafIdxEnd {
			upper = nodes[2*i+2].upper
			upperInc = nodes[2*i+2].upperInc
		}
		nodes[i] = &segmentNode{valueRange: &valueRange{lower: lower, upper: upper, lowerInc: lowerInc, upperInc: upperInc}}
	}
	buildTree(0)

	// insert the intervals into the tree
	var insIntv func(i int, rg *idxRange)
	insIntv = func(i int, rg *idxRange) {
		// rg contains nodes[i]
		if contains(rg.valueRange, nodes[i].valueRange) {
			nodes[i].idx = add(nodes[i].idx, rg.index)
			return
		}
		if 2*i+1 < leafIdxEnd && intersects(rg.valueRange, nodes[2*i+1].valueRange) {
			insIntv(2*i+1, rg)
		}
		if 2*i+2 < leafIdxEnd && intersects(rg.valueRange, nodes[2*i+2].valueRange) {
			insIntv(2*i+2, rg)
		}
	}
	for _, rg := range rgs {
		insIntv(0, rg)
	}

	return &segmentTree{
		nodes: nodes,
	}
}

func (c *segmentTree) find(val T) Bitmap {
	if c == nil {
		return nil
	}
	return c.findAt(0, val)
}

func (c *segmentTree) findAt(i int, v T) Bitmap {
	if i >= len(c.nodes) {
		return nil
	}
	// v is not contained
	if !containsVal(c.nodes[i].valueRange, v) {
		return nil
	}
	return or(c.nodes[i].idx, or(c.findAt(2*i+1, v), c.findAt(2*i+2, v)))
}