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);
 */