LeetCode: 307. Range Sum Query - Mutable

LeetCode: 307. Range Sum Query - Mutable

题目描述

Given an integer array nums, find the sum of the elements between indices ​​i​​​ and ​​j (i ≤ j)​​, inclusive.

The ​​update(i, val)​​​ function modifies nums by updating the element at index ​​i​​​ to ​​val​​.

Example:

Given nums = [1, 3, 5]

sumRange(0, 2) -> 9
update(1, 2)
sumRange(0, 2) -> 8

Note:

1. The array is only modifiable by the ​​update​​ function.
2. You may assume the number of calls to ​​update​​​ and ​​sumRange​​ function is distributed evenly.

解题思路 —— 线段树(segment tree)

1. 将给定数组扩展成满二叉树的叶子，不足的部分用 ​​0​​ 补足
2. 每次更新数组时，循环更新它的祖先节点
3. 每次查询区间时，从根节点递归向下寻找满足在指定区间内的节点

AC 代码

type NumArray struct {
    tree []int
}

func Constructor(nums []int) NumArray {
    var arr NumArray
    
    if len(nums) == 0 { return arr }
    
    // 形成满二叉树
    for len(nums)&(len(nums)-1) != 0 { 
        nums = append(nums, 0) 
    }
    
    arr.tree = make([]int, len(nums)*2)
    arr.tree[0] = -1 
    
    // 生成线段树
    for i := len(nums); i >= 1; i /= 2 {
        for j := i; j < i*2; j++{
            if j >= len(nums) {
                arr.tree[j] = nums[j-len(nums)]
            } else {
                arr.tree[j] = arr.tree[j*2] + arr.tree[j*2+1]
            }
        }
    }
    
    return arr
}


func (this *NumArray) Update(i int, val int)  {
    // 更新线段树
    for j := len(this.tree)/2 + i; j >= 1; j /= 2 {
        if j == len(this.tree)/2 + i {
            this.tree[j] = val
        } else {
            this.tree[j] = this.tree[j*2] + this.tree[j*2+1];
        }
    }

}

func min(m, n int) int {
    if m > n {
        return n
    } else {
        return m
    }
}

func max(m, n int) int {
    return (m + n) - min(m, n)
}

func (this *NumArray) SumRangeRef(i int, j int, idx int, beg int, end int) int { 
    
    // 区间 [i, j] 和 区间 [beg, end] 不相交
    if beg > j || end < i || end < beg {
        return 0
    }
    
    
    // 区间 [beg, end] 包含于区间 [i, j]
    if beg >= i && end <= j {
        return this.tree[idx]
    }
    
    ans := 0
    mid := (beg + end) / 2
    
    // 区间 [beg, mid] 和 区间 [i, j] 的交集
    l1 := max(beg, i)
    r1 := min(j, mid)
    if l1 <= r1 {
        ans += this.SumRangeRef(l1, r1, idx*2, beg, mid)
    }
    
    // 区间 [mid+1, end] 和 区间 [i, j] 的交集
    l2 := max(mid+1, i)
    r2 := min(j, end)
    if l2 <= r2 {
        ans += this.SumRangeRef(l2, r2, idx*2+1, mid+1, end)
    }
    
    return ans
}

func (this *NumArray) SumRange(i int, j int) int {
    return this.SumRangeRef(i+1, j+1, 1, 1, len(this.tree)/2)
}


/**
 * Your NumArray object will be instantiated and called as such:
 * obj := Constructor(nums);
 * obj.Update(i,val);
 * param_2 := obj.SumRange(i,j);
 */