数据结构板子 in Golang

线段树

区间求和、区间加

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
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
// https://www.luogu.com.cn/problem/P3372
package main

import (
"bufio"
"fmt"
"io"
"os"
)

type SegmentTree struct {
leftChild *SegmentTree
rightChild *SegmentTree
begin int
endInclusive int
mid int
inc int
sum int
}

func NewSegmentTree(begin int, endInclusive int) *SegmentTree {
root := newNode(begin, endInclusive)
root.growDown()
return root
}

func newNode(begin int, endInclusive int) *SegmentTree {
return &SegmentTree{
leftChild: nil,
rightChild: nil,
begin: begin,
endInclusive: endInclusive,
mid: (begin + endInclusive) >> 1,
inc: 0,
sum: 0,
}
}

func (node *SegmentTree) growDown() {
if node.endInclusive == node.begin {
return
}

node.leftChild = newNode(node.begin, node.mid)
node.rightChild = newNode(node.mid+1, node.endInclusive)

node.leftChild.growDown()
node.rightChild.growDown()
}

func (node *SegmentTree) SumUp(qBegin int, qEndInclusive int) int {
if node.begin == qBegin && node.endInclusive == qEndInclusive {
return node.sum
}

ans := node.inc * (qEndInclusive - qBegin + 1)
if qBegin <= node.mid {
if node.mid < qEndInclusive {
ans += node.leftChild.SumUp(qBegin, node.mid)
} else {
ans += node.leftChild.SumUp(qBegin, qEndInclusive)
}
}
if node.mid+1 <= qEndInclusive {
if node.mid+1 > qBegin {
ans += node.rightChild.SumUp(node.mid+1, qEndInclusive)
} else {
ans += node.rightChild.SumUp(qBegin, qEndInclusive)
}
}
return ans
}

func (node *SegmentTree) AddEach(qBegin int, qEndInclusive int, value int) {
node.sum += (qEndInclusive - qBegin + 1) * value
if node.begin == qBegin && node.endInclusive == qEndInclusive {
node.inc += value
return
}

if qBegin <= node.mid {
if node.mid < qEndInclusive {
node.leftChild.AddEach(qBegin, node.mid, value)
} else {
node.leftChild.AddEach(qBegin, qEndInclusive, value)
}
}
if node.mid+1 <= qEndInclusive {
if node.mid+1 > qBegin {
node.rightChild.AddEach(node.mid+1, qEndInclusive, value)
} else {
node.rightChild.AddEach(qBegin, qEndInclusive, value)
}
}
}

func run(_r io.Reader, _w io.Writer) {
in := bufio.NewReader(_r)
out := bufio.NewWriter(_w)
defer out.Flush()

var n, q, op, l, r, val int

fmt.Fscan(in, &n, &q)
t := NewSegmentTree(1, n)

for i := 1; i <= n; i++ {
fmt.Fscan(in, &val)
t.AddEach(i, i, val)
}

for ; q > 0; q-- {
fmt.Fscan(in, &op, &l, &r)
if op == 1 {
fmt.Fscan(in, &val)
t.AddEach(l, r, val)
} else {
fmt.Fprintln(out, t.SumUp(l, r))
}
}

}
func main() {
run(os.Stdin, os.Stdout)
}
Author: ssttkkl
Link: https://ssttkkl.github.io/uncategorized/2020/07/%E6%95%B0%E6%8D%AE%E7%BB%93%E6%9E%84%E6%9D%BF%E5%AD%90-in-Golang/
Copyright Notice: All articles in this blog are licensed under CC BY-NC-SA 4.0 unless stating additionally.