ADD: BIT

app.go

@@ -1,6 +1,6 @@
 /*
  * @Author: NorthCity1984
- * @LastEditTime: 2022-04-26 15:55:37
+ * @LastEditTime: 2022-04-26 17:05:06
  * @Description:
  * @Website: https://grimoire.cn
  * Copyright (c) NorthCity1984 All rights reserved.
@@ -10,26 +10,18 @@ package main
 import (
 	"fmt"
 
-	"gitee.com/NorthCityChen/stl-go/heap"
+	"gitee.com/NorthCityChen/stl-go/bit"
 )
 
-// import "container/heap"
-
-var a = []int{1, 2, 65, 4, 5}
+var a = []int{1, 2, 65, 4, 33, 5}
 
 func main() {
-	// <: 小根堆 >: 大根堆
-	f := func(t1, t2 int) bool { return t1 > t2 }
-	h := heap.Init(a, f)
-	fmt.Println(h)
-
-	fmt.Println(h)
-
-	fmt.Println(h.Pop())
-	fmt.Println(h.Pop())
-	fmt.Println(h.Pop())
-	fmt.Println(h.Pop())
-	fmt.Println(h.Pop())
-	fmt.Println(h.Pop())
-	fmt.Println(h.Pop())
+	ans := a[0]
+	for i := len(a) - 1; i >= 1; i-- {
+		a[i] = a[i] - a[i-1]
+	}
+	b := bit.Init(a)
+	fmt.Println(b)
+	ans = b.Ask(6)
+	fmt.Println(ans)
 }

bit/bit.go

@@ -0,0 +1,55 @@
+/*
+ * @Author: NorthCity1984
+ * @LastEditTime: 2022-04-26 17:08:18
+ * @Description: Binary Indexed Trees
+ * @Website: https://grimoire.cn
+ * Copyright (c) NorthCity1984 All rights reserved.
+ */
+
+package bit
+
+type Number interface {
+	int | int64 | float32 | float64
+}
+
+type bit[T Number] struct {
+	array []T
+	brray []T
+	len   int
+}
+
+func lowbit(x int) int {
+	return x & -x
+}
+
+// 初始化一个树状数组
+func Init[T Number](array []T) bit[T] {
+	ret := bit[T]{
+		array: append([]T{0}, array...),
+		brray: make([]T, len(array)+5),
+		len:   len(array),
+	}
+
+	for i := 1; i < len(ret.array); i++ {
+		ret.Add(i, ret.array[i])
+	}
+	return ret
+}
+
+// 在第x个数的位置+k
+func (b *bit[T]) Add(x int, k T) {
+	for x <= b.len {
+		b.brray[x] = b.brray[x] + k
+		x = x + lowbit(x) // 父节点
+	}
+}
+
+// 求第一个数到第x个数的和
+func (b *bit[T]) Ask(x int) T {
+	var ans T = 0
+	for x >= 1 {
+		ans = ans + b.brray[x]
+		x = x - lowbit(x)
+	}
+	return ans
+}

heap/heap.go

@@ -1,6 +1,6 @@
 /*
  * @Author: NorthCity1984
- * @LastEditTime: 2022-04-26 15:55:32
+ * @LastEditTime: 2022-04-26 16:06:24
  * @Description:
  * @Website: https://grimoire.cn
  * Copyright (c) NorthCity1984 All rights reserved.
@@ -11,29 +11,20 @@ type Number interface {
 	int | int64 | float32 | float64
 }
 
-type Heap[T Number] struct {
+type heap[T Number] struct {
 	array []T
 	len   int
 	cmp   func(T, T) bool
 }
 
-func (h *Heap[T]) up(x int) {
-	// if h.isMax {
-
+func (h *heap[T]) up(x int) {
 	for x > 1 && h.cmp(h.array[x], h.array[x/2]) {
 		h.array[x], h.array[x/2] = h.array[x/2], h.array[x]
 		x /= 2
 	}
-	// } else {
-	// for x > 1 && h.array[x] < h.array[x/2] {
-	// 	h.array[x], h.array[x/2] = h.array[x/2], h.array[x]
-	// 	x /= 2
-	// }
-	// }
-
 }
 
-func (h *Heap[T]) down(x int) {
+func (h *heap[T]) down(x int) {
 	for x*2 <= h.len {
 		t := x * 2
 		if t+1 <= h.len && h.cmp(h.array[t+1], h.array[t]) {
@@ -47,13 +38,13 @@ func (h *Heap[T]) down(x int) {
 	}
 }
 
-func (h *Heap[T]) Push(v T) {
+func (h *heap[T]) Push(v T) {
 	h.array = append(h.array, v)
 	h.len++
 	h.up(h.len)
 }
 
-func (h *Heap[T]) Pop() T {
+func (h *heap[T]) Pop() T {
 	if h.len == 0 {
 		return 0
 	}
@@ -61,18 +52,16 @@ func (h *Heap[T]) Pop() T {
 	h.array[1], h.array[h.len] = h.array[h.len], h.array[1]
 	h.array = h.array[:h.len]
 	h.len--
-	// fmt.Println(h)
 	h.down(1)
-
 	return ret
 }
 
-func (h *Heap[T]) Top() T {
+func (h *heap[T]) Top() T {
 	return h.array[1]
 }
 
-func Init[T Number](array []T, cmp func(T, T) bool) Heap[T] {
-	h := Heap[T]{
+func Init[T Number](array []T, cmp func(T, T) bool) heap[T] {
+	h := heap[T]{
 		array: append([]T{0}, array...),
 		len:   len(array),
 		cmp:   cmp,
@@ -80,8 +69,5 @@ func Init[T Number](array []T, cmp func(T, T) bool) Heap[T] {
 	for i := h.len; i >= 1; i-- {
 		h.down(i)
 	}
-	// for i := 1; i <= h.len; i++ {
-	// 	h.up(i)
-	// }
 	return h
 }