307 更快的解法

Author: mlkr

Algorithms/0307.range-sum-query-mutable/range-sum-query-mutable.go

@@ -1,27 +1,80 @@
 package problem0307
 
-// NumArray 是 nums 的和的切片
+// https://leetcode.com/problems/range-sum-query-mutable/discuss/75724/17-ms-Java-solution-with-segment-tree
+type SegmentTreeNode struct {
+	start, end, sum int
+	left, right     *SegmentTreeNode
+}
+
+func (node *SegmentTreeNode) update(i, val int) {
+	if node.start == node.end {
+		node.sum = val
+		return
+	}
+
+	mid := node.start + (node.end-node.start)/2
+	if mid >= i {
+		node.left.update(i, val)
+	} else {
+		node.right.update(i, val)
+	}
+
+	node.sum = node.left.sum + node.right.sum
+}
+
+func (node *SegmentTreeNode) sumRange(i int, j int) int {
+	if node.start == i && node.end == j {
+		return node.sum
+	}
+
+	mid := node.start + (node.end-node.start)/2
+	if mid >= j {
+		return node.left.sumRange(i, j)
+	} else if mid < i {
+		return node.right.sumRange(i, j)
+	} else {
+		return node.left.sumRange(i, mid) + node.right.sumRange(mid+1, j)
+	}
+}
+
+func buildTree(nums []int, start, end int) *SegmentTreeNode {
+	if start > end {
+		return nil
+	}
+
+	node := &SegmentTreeNode{
+		start: start,
+		end:   end,
+	}
+
+	if start == end {
+		node.sum = nums[start]
+	} else {
+		mid := start + (end-start)/2
+		node.left = buildTree(nums, start, mid)
+		node.right = buildTree(nums, mid+1, end)
+		node.sum = node.left.sum + node.right.sum
+	}
+
+	return node
+}
+
 type NumArray struct {
-	nums []int
+	root *SegmentTreeNode
 }
 
-// Constructor 返回 NumArray
 func Constructor(nums []int) NumArray {
-	return NumArray{nums: nums}
+	return NumArray{
+		buildTree(nums, 0, len(nums)-1),
+	}
 }
 
-// Update 更新 nums
-func (na *NumArray) Update(i int, v int) {
-	na.nums[i] = v
+func (this *NumArray) Update(i int, val int) {
+	this.root.update(i, val)
 }
 
-// SumRange 返回 sum(nums[i:j+1])
-func (na *NumArray) SumRange(i int, j int) int {
-	res := 0
-	for k := i; k <= j; k++ {
-		res += na.nums[k]
-	}
-	return res
+func (this *NumArray) SumRange(i int, j int) int {
+	return this.root.sumRange(i, j)
 }
 
 /**

Algorithms/0307.range-sum-query-mutable/range-sum-query-mutable_test.go

@@ -6,30 +6,30 @@ import (
 	"github.com/stretchr/testify/assert"
 )
 
-func Test_SumRange(t *testing.T) {
-	ast := assert.New(t)
-
-	nums := []int{1, 3, 5}
-
-	na := Constructor(nums)
-
-	ast.Equal(9, na.SumRange(0, 2), "update 前，SumRange(0, 2)")
-
-	na.Update(1, 2)
+var nums = []int{1, 3, 5}
+var numsnil = []int{}
 
-	ast.Equal(8, na.SumRange(0, 2), "update 后，SumRange(0, 2)")
-}
-
-func Test_SumRange_2(t *testing.T) {
+func Test_NumArray(t *testing.T) {
 	ast := assert.New(t)
+	numArray := Constructor(nums)
 
-	nums := []int{-1}
+	ast.Equal(9, numArray.SumRange(0, 2))
+	numArray.Update(1, 2)
+	ast.Equal(8, numArray.SumRange(0, 2))
+	ast.Equal(7, numArray.SumRange(1, 2))
+	ast.Equal(1, numArray.SumRange(0, 0))
 
-	na := Constructor(nums)
+	numArray1 := Constructor(numsnil)
+	ast.Nil(numArray1.root)
 
-	ast.Equal(-1, na.SumRange(0, 0), "update 前，SumRange(0, 0)")
-
-	na.Update(0, 1)
+}
 
-	ast.Equal(1, na.SumRange(0, 0), "update 后，SumRange(0, 2)")
+func Benchmark_NumArray(b *testing.B) {
+	for i := 0; i < b.N; i++ {
+		numArray := Constructor(nums)
+		numArray.SumRange(0, 2)
+		numArray.Update(1, 2)
+		numArray.SumRange(0, 2)
+		numArray.SumRange(1, 2)
+	}
 }