Range queries: Fenwick & segment trees

package main
 
import (
	"fmt"
	"slices"
)
 
// Fenwick supports point updates and prefix sums in O(log n). Indexes are 0-based outside, 1-based inside.
type Fenwick struct{ tree []int }
 
func NewFenwick(n int) *Fenwick { return &Fenwick{make([]int, n+1)} }
 
func (f *Fenwick) Add(i, delta int) {
	for i++; i < len(f.tree); i += i & -i {
		f.tree[i] += delta
	}
}
 
// Prefix returns the sum of elements [0, i).
func (f *Fenwick) Prefix(i int) int {
	s := 0
	for ; i > 0; i -= i & -i {
		s += f.tree[i]
	}
	return s
}
 
func (f *Fenwick) Range(l, r int) int { return f.Prefix(r+1) - f.Prefix(l) } // inclusive
 
// countSmallerAfter: for each element, how many smaller elements are to its right?
func countSmallerAfter(nums []int) []int {
	// Coordinate compression: map values to ranks 0..m-1 so the tree stays small.
	sorted := slices.Clone(nums)
	slices.Sort(sorted)
	sorted = slices.Compact(sorted)
	f := NewFenwick(len(sorted))
	out := make([]int, len(nums))
	for i := len(nums) - 1; i >= 0; i-- {
		rank, _ := slices.BinarySearch(sorted, nums[i])
		out[i] = f.Prefix(rank) // how many seen values have a smaller rank
		f.Add(rank, 1)
	}
	return out
}
 
func main() {
	nums := []int{5, 8, 6, 3, 2, 7, 2, 6}
	f := NewFenwick(len(nums))
	for i, v := range nums {
		f.Add(i, v)
	}
	fmt.Println("sum[2..6] =", f.Range(2, 6))
	f.Add(4, 10) // nums[4] += 10
	fmt.Println("after nums[4] += 10, sum[2..6] =", f.Range(2, 6))
	fmt.Println("internal tree:", f.tree[1:])
 
	fmt.Println("smaller after self:", countSmallerAfter([]int{5, 2, 6, 1}))
}

sum[2..6] = 20
after nums[4] += 10, sum[2..6] = 30
internal tree: [5 13 6 22 12 19 2 49]
smaller after self: [2 1 1 0]

package main
 
import (
	"fmt"
	"math"
)
 
// MinTree answers range-minimum queries with point updates.
type MinTree struct {
	n int
	t []int
}
 
func NewMinTree(a []int) *MinTree {
	n := len(a)
	t := make([]int, 2*n)
	copy(t[n:], a) // leaves
	for i := n - 1; i > 0; i-- {
		t[i] = min(t[2*i], t[2*i+1]) // each parent combines its two children
	}
	return &MinTree{n, t}
}
 
func (m *MinTree) Set(i, v int) {
	i += m.n
	m.t[i] = v
	for i /= 2; i > 0; i /= 2 {
		m.t[i] = min(m.t[2*i], m.t[2*i+1])
	}
}
 
// Min returns the minimum of a[l..r], inclusive.
func (m *MinTree) Min(l, r int) int {
	res := math.MaxInt
	for l, r = l+m.n, r+m.n+1; l < r; l, r = l/2, r/2 {
		if l&1 == 1 { // l is a right child: take it alone, then move right
			res = min(res, m.t[l])
			l++
		}
		if r&1 == 1 { // r is exclusive; if it is a right child, its left sibling is inside
			r--
			res = min(res, m.t[r])
		}
	}
	return res
}
 
func main() {
	a := []int{5, 8, 6, 3, 2, 7, 2, 6}
	m := NewMinTree(a)
	fmt.Println(m.Min(0, 2), m.Min(2, 6), m.Min(5, 7), m.Min(0, 7))
	m.Set(4, 9) // a[4] = 9
	fmt.Println("after a[4] = 9:", m.Min(2, 6), m.Min(3, 4))
 
	// Check against a brute-force scan on every range.
	ok := true
	for l := range a {
		for r := l; r < len(a); r++ {
			want := math.MaxInt
			for _, v := range []int{5, 8, 6, 3, 9, 7, 2, 6}[l : r+1] {
				want = min(want, v)
			}
			ok = ok && m.Min(l, r) == want
		}
	}
	fmt.Println("all 36 ranges match brute force:", ok)
}

5 2 2 2
after a[4] = 9: 2 3
all 36 ranges match brute force: true

range [2, 8], length 7, k = 2 (block length 4):
  block 1: [2, 5]      block 2: [5, 8]      overlap at 5 is harmless for min

package main
 
import (
	"fmt"
	"math/bits"
)
 
type Sparse struct{ sp [][]int }
 
func NewSparse(a []int) *Sparse {
	levels := bits.Len(uint(len(a)))
	sp := make([][]int, levels)
	sp[0] = append([]int(nil), a...)
	for k := 1; k < levels; k++ {
		half := 1 << (k - 1)
		sp[k] = make([]int, len(a)-(1<<k)+1)
		for i := range sp[k] {
			sp[k][i] = min(sp[k-1][i], sp[k-1][i+half]) // two halves of length 2^(k-1)
		}
	}
	return &Sparse{sp}
}
 
// Min of a[l..r] in O(1): two overlapping blocks of length 2^k.
func (s *Sparse) Min(l, r int) int {
	k := bits.Len(uint(r-l+1)) - 1
	return min(s.sp[k][l], s.sp[k][r-(1<<k)+1])
}
 
func main() {
	s := NewSparse([]int{5, 8, 6, 3, 2, 7, 2, 6, 4, 1})
	for k, row := range s.sp {
		fmt.Printf("blocks of %d: %v\n", 1<<k, row)
	}
	fmt.Println(s.Min(2, 8), s.Min(0, 2), s.Min(7, 9))
}

blocks of 1: [5 8 6 3 2 7 2 6 4 1]
blocks of 2: [5 6 3 2 2 2 2 4 1]
blocks of 4: [3 2 2 2 2 2 1]
blocks of 8: [2 2 1]
2 5 1

package main
 
import "fmt"
 
const LOG = 4 // enough for trees with up to 16 levels
 
func main() {
	//            0
	//          /   \
	//         1     2
	//        / \     \
	//       3   4     5
	//      /         / \
	//     6         7   8
	parent := []int{0, 0, 0, 1, 1, 2, 3, 5, 5} // the root is its own parent
	n := len(parent)
	depth := make([]int, n)
	for v := 1; v < n; v++ {
		depth[v] = depth[parent[v]] + 1 // parents come before children here
	}
	var up [LOG][]int
	up[0] = parent
	for k := 1; k < LOG; k++ {
		up[k] = make([]int, n)
		for v := range n {
			up[k][v] = up[k-1][up[k-1][v]] // two jumps of 2^(k-1)
		}
	}
 
	kth := func(v, k int) int {
		for j := 0; j < LOG; j++ {
			if k>>j&1 == 1 {
				v = up[j][v]
			}
		}
		return v
	}
	lca := func(u, v int) int {
		if depth[u] < depth[v] {
			u, v = v, u
		}
		u = kth(u, depth[u]-depth[v]) // same depth
		if u == v {
			return u
		}
		for j := LOG - 1; j >= 0; j-- {
			if up[j][u] != up[j][v] {
				u, v = up[j][u], up[j][v]
			}
		}
		return parent[u]
	}
 
	fmt.Println("2nd ancestor of 6:", kth(6, 2), " 3rd ancestor of 7:", kth(7, 3))
	for _, q := range [][2]int{{6, 4}, {6, 7}, {7, 8}, {3, 6}} {
		a := lca(q[0], q[1])
		fmt.Printf("lca(%d,%d) = %d, distance %d\n", q[0], q[1], a, depth[q[0]]+depth[q[1]]-2*depth[a])
	}
}

2nd ancestor of 6: 1  3rd ancestor of 7: 0
lca(6,4) = 1, distance 3
lca(6,7) = 0, distance 6
lca(7,8) = 5, distance 2
lca(3,6) = 3, distance 1