1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47

// @Title: 区域和检索 - 数组可修改 (Range Sum Query - Mutable)
// @Author: 15816537946@163.com
// @Date: 2022-04-04 10:15:33
// @Runtime: 532 ms
// @Memory: 22.5 MB
type NumArray struct {
    nums, sums []int // sums[i] 表示第 i 个块的元素和
    size       int   // 块的大小
}

func Constructor(nums []int) NumArray {
    n := len(nums)
    size := int(math.Sqrt(float64(n)))
    sums := make([]int, (n+size-1)/size) // n/size 向上取整
    for i, num := range nums {
        sums[i/size] += num
    }
    return NumArray{nums, sums, size}
}

func (na *NumArray) Update(index, val int) {
    na.sums[index/na.size] += val - na.nums[index]
    na.nums[index] = val
}

func (na *NumArray) SumRange(left, right int) (ans int) {
    size := na.size
    b1, b2 := left/size, right/size
    if b1 == b2 { // 区间 [left, right] 在同一块中
        for i := left; i <= right; i++ {
            ans += na.nums[i]
        }
        return
    }
    for i := left; i < (b1+1)*size; i++ {
        ans += na.nums[i]
    }
    for i := b1 + 1; i < b2; i++ {
        ans += na.sums[i]
    }
    for i := b2 * size; i <= right; i++ {
        ans += na.nums[i]
    }
    return
}