結果

問題 No.1038 TreeAddQuery
ユーザー 草苺奶昔草苺奶昔
提出日時 2023-12-02 15:28:43
言語 Go
(1.23.4)
結果
AC  
実行時間 2,254 ms / 4,000 ms
コード長 11,707 bytes
コンパイル時間 15,207 ms
コンパイル使用メモリ 220,272 KB
実行使用メモリ 286,916 KB
最終ジャッジ日時 2024-09-26 18:47:03
合計ジャッジ時間 48,641 ms
ジャッジサーバーID
(参考情報)
judge4 / judge1
このコードへのチャレンジ
(要ログイン)
ファイルパターン 結果
sample AC * 3
other AC * 24
権限があれば一括ダウンロードができます

ソースコード

diff #
プレゼンテーションモードにする

package main
import (
"bufio"
"fmt"
"os"
"strings"
)
func main() {
// VertexGetRangeContourAdd()
Yuki1038()
}
// https://judge.yosupo.jp/problem/vertex_add_range_contour_sum_on_tree
// q
// 0 root x : rootx ()
// 1 root floor higher: root[floor,higher) ()
// n<=1e5 q<=2e5
func VertexAddRangeContourSum() {
in := bufio.NewReader(os.Stdin)
out := bufio.NewWriter(os.Stdout)
defer out.Flush()
var n, q int
fmt.Fscan(in, &n, &q)
weights := make([]int, n)
for i := range weights {
fmt.Fscan(in, &weights[i])
}
tree := make([][]int, n)
for i := 0; i < n-1; i++ {
var a, b int
fmt.Fscan(in, &a, &b)
tree[a] = append(tree[a], b)
tree[b] = append(tree[b], a)
}
C := NewContourQueryRange(n, tree)
data := make([]int, C.Size())
for i := 0; i < n; i++ {
C.EnumeratePoint(i, func(pos int) {
data[pos] += weights[i]
})
}
bit := NewBitArrayFrom(len(data), func(i int) int { return data[i] })
for i := 0; i < q; i++ {
var kind int
fmt.Fscan(in, &kind)
if kind == 0 {
var root, x int
fmt.Fscan(in, &root, &x)
weights[root] += x
C.EnumeratePoint(root, func(pos int) {
bit.Add(pos, x)
})
} else {
var root, floor, higher int
fmt.Fscan(in, &root, &floor, &higher)
res := 0
if floor <= 0 && 0 < higher {
res += weights[root]
}
C.EnumerateRange(root, floor, higher, func(start, end int) {
res += bit.QueryRange(start, end)
})
fmt.Fprintln(out, res)
}
}
}
// https://judge.yosupo.jp/problem/vertex_get_range_contour_add_on_tree
// q
// 0 root floor higher x: root[floor,higher)x ()
// 1 root : root ()
// n<=1e5 q<=2e5
func VertexGetRangeContourAdd() {
in := bufio.NewReader(os.Stdin)
out := bufio.NewWriter(os.Stdout)
defer out.Flush()
var n, q int
fmt.Fscan(in, &n, &q)
values := make([]int, n)
for i := range values {
fmt.Fscan(in, &values[i])
}
tree := make([][]int, n)
for i := 0; i < n-1; i++ {
var a, b int
fmt.Fscan(in, &a, &b)
tree[a] = append(tree[a], b)
tree[b] = append(tree[b], a)
}
C := NewContourQueryRange(n, tree)
bit := NewBitArray(C.Size() + 1)
add := func(root, floor, higher, x int) {
if floor <= 0 && 0 < higher {
values[root] += x
}
C.EnumerateRange(root, floor, higher, func(start, end int) {
bit.Add(start, x)
bit.Add(end, -x)
})
}
query := func(root int) int {
res := values[root]
C.EnumeratePoint(root, func(pos int) {
res += bit.QueryPrefix(pos + 1)
})
return res
}
for i := 0; i < q; i++ {
var kind int
fmt.Fscan(in, &kind)
if kind == 0 {
var root, floor, higher, x int
fmt.Fscan(in, &root, &floor, &higher, &x)
add(root, floor, higher, x)
} else {
var root int
fmt.Fscan(in, &root)
fmt.Fprintln(out, query(root))
}
}
}
// https://yukicoder.me/problems/no/1038
// 0q.
// xx[0,y]z.
func Yuki1038() {
in, out := bufio.NewReader(os.Stdin), bufio.NewWriter(os.Stdout)
defer out.Flush()
var n, q int
fmt.Fscan(in, &n, &q)
tree := make([][]int, n)
for i := 0; i < n-1; i++ {
var a, b int
fmt.Fscan(in, &a, &b)
a--
b--
tree[a] = append(tree[a], b)
tree[b] = append(tree[b], a)
}
weights := make([]int, n)
C := NewContourQueryRange(n, tree)
bit := NewBitArray(C.Size() + 1)
for i := 0; i < q; i++ {
var node, dist, delta int
fmt.Fscan(in, &node, &dist, &delta)
node--
dist++
res := weights[node]
C.EnumeratePoint(node, func(pos int) {
res += bit.QueryPrefix(pos + 1)
})
fmt.Fprintln(out, res)
weights[node] += delta
C.EnumerateRange(node, 0, dist, func(start, end int) {
bit.Add(start, delta)
bit.Add(end, -delta)
})
}
}
// !0.
type ContourQueryRange struct {
_n int
_v []int
_comp []int
_dep []int
_infoIdx []int
_infoIndptr []int
_compRange []int
}
func NewContourQueryRange(n int, graph [][]int) *ContourQueryRange {
p := 0
compRange := []int{0}
V := []int{}
comp := []int{}
dep := []int{}
f := func(par []int, vs []int, color []int) {
n := len(par)
dist := make([]int, n)
A := make([]int, 0)
B := make([]int, 0)
for v := 1; v < n; v++ {
dist[v] = dist[par[v]] + 1
}
for v := 0; v < n; v++ {
if color[v] == 0 {
A = append(A, v)
}
if color[v] == 1 {
B = append(B, v)
}
}
if len(A) == 0 || len(B) == 0 {
return
}
mxA := 0
mxB := 0
for _, v := range A {
V = append(V, vs[v])
comp = append(comp, p)
dep = append(dep, dist[v])
mxA = max(mxA, dist[v])
}
compRange = append(compRange, compRange[len(compRange)-1]+mxA+1)
p++
for _, v := range B {
V = append(V, vs[v])
comp = append(comp, p)
dep = append(dep, dist[v])
mxB = max(mxB, dist[v])
}
compRange = append(compRange, compRange[len(compRange)-1]+mxB+1)
p++
}
centroidDecomposition(n, graph, f)
infoIndptr := make([]int, n+1)
for _, v := range V {
infoIndptr[v+1]++
}
for v := 0; v < n; v++ {
infoIndptr[v+1] += infoIndptr[v]
}
counter := append([]int{}, infoIndptr...)
infoIdx := make([]int, infoIndptr[len(infoIndptr)-1])
for i := 0; i < len(V); i++ {
infoIdx[counter[V[i]]] = i
counter[V[i]]++
}
return &ContourQueryRange{
_n: n,
_v: V,
_comp: comp,
_dep: dep,
_infoIdx: infoIdx,
_infoIndptr: infoIndptr,
_compRange: compRange,
}
}
func (cqr *ContourQueryRange) Size() int {
return cqr._compRange[len(cqr._compRange)-1]
}
func (cqr *ContourQueryRange) EnumerateRange(node int, start int, end int, f func(int, int)) {
for k := cqr._infoIndptr[node]; k < cqr._infoIndptr[node+1]; k++ {
idx := cqr._infoIdx[k]
p := cqr._comp[idx] ^ 1
lo := start - cqr._dep[idx]
hi := end - cqr._dep[idx]
L := cqr._compRange[p]
R := cqr._compRange[p+1]
n := R - L
lo = max(lo, 0)
hi = min(hi, n)
if lo < hi {
f(L+lo, L+hi)
}
}
}
func (cqr *ContourQueryRange) EnumeratePoint(v int, f func(int)) {
for k := cqr._infoIndptr[v]; k < cqr._infoIndptr[v+1]; k++ {
idx := cqr._infoIdx[k]
p := cqr._comp[idx]
f(cqr._compRange[p] + cqr._dep[idx])
}
}
func centroidDecomposition(n int, g [][]int, f func([]int, []int, []int)) {
if n == 1 {
return
}
V := make([]int, n)
par := make([]int, n)
for i := range par {
par[i] = -1
}
l := 0
r := 0
V[r] = 0
r++
for l < r {
v := V[l]
l++
for _, next := range g[v] {
if next != par[v] {
V[r] = next
par[next] = v
r++
}
}
}
if r != n {
panic("r should be equal to n")
}
newIdx := make([]int, n)
for i := 0; i < n; i++ {
newIdx[V[i]] = i
}
tmp := make([]int, n)
for i := 0; i < n; i++ {
tmp[i] = -1
}
for i := 1; i < n; i++ {
j := par[i]
tmp[newIdx[i]] = newIdx[j]
}
par = tmp
real := make([]int, n)
for i := range real {
real[i] = 1
}
centroidDecomposition2DFS(par, V, real, f)
}
func centroidDecomposition2DFS(par []int, vs []int, real []int, f func([]int, []int, []int)) {
N := len(par)
if N <= 1 {
panic("N should be greater than or equal to 2")
}
if N == 2 {
if real[0] != 0 && real[1] != 0 {
color := []int{0, 1}
f(par, vs, color)
}
return
}
c := -1
sz := make([]int, N)
for i := range sz {
sz[i] = 1
}
for i := N - 1; i >= 0; i-- {
if sz[i] >= (N+1)>>1 {
c = i
break
}
sz[par[i]] += sz[i]
}
color := make([]int, N)
for i := range color {
color[i] = -1
}
take := 0
ord := make([]int, N)
for i := range ord {
ord[i] = -1
}
ord[c] = 0
p := 1
for v := 1; v < N; v++ {
if par[v] == c && take+sz[v] <= (N-1)/2 {
color[v] = 0
ord[v] = p
p++
take += sz[v]
}
}
for i := 1; i < N; i++ {
if color[par[i]] == 0 {
color[i] = 0
ord[i] = p
p++
}
}
n0 := p - 1
for a := par[c]; a != -1; a = par[a] {
color[a] = 1
ord[a] = p
p++
}
for i := 0; i < N; i++ {
if i != c && color[i] == -1 {
color[i] = 1
ord[i] = p
p++
}
}
if p != N {
panic("p should be equal to N")
}
n1 := N - 1 - n0
par0 := make([]int, n0+1)
for i := range par0 {
par0[i] = -1
}
par1 := make([]int, n1+1)
for i := range par1 {
par1[i] = -1
}
par2 := make([]int, N)
for i := range par2 {
par2[i] = -1
}
V0 := make([]int, n0+1)
V1 := make([]int, n1+1)
V2 := make([]int, N)
rea0 := make([]int, n0+1)
rea1 := make([]int, n1+1)
rea2 := make([]int, N)
for v := 0; v < N; v++ {
i := ord[v]
V2[i] = vs[v]
rea2[i] = real[v]
if color[v] != 1 {
V0[i] = vs[v]
rea0[i] = real[v]
}
if color[v] != 0 {
V1[max(i-n0, 0)] = vs[v]
rea1[max(i-n0, 0)] = real[v]
}
}
for v := 1; v < N; v++ {
a := ord[v]
b := ord[par[v]]
if a > b {
a, b = b, a
}
par2[b] = a
if color[v] != 1 && color[par[v]] != 1 {
par0[b] = a
}
if color[v] != 0 && color[par[v]] != 0 {
par1[max(b-n0, 0)] = max(a-n0, 0)
}
}
if real[c] != 0 {
color = make([]int, N)
for i := 0; i < N; i++ {
color[i] = -1
}
color[0] = 0
for i := 1; i < N; i++ {
if rea2[i] != 0 {
color[i] = 1
} else {
color[i] = -1
}
}
f(par2, V2, color)
rea0[0] = 0
rea1[0] = 0
rea2[0] = 0
}
color = make([]int, N)
for i := 0; i < N; i++ {
color[i] = -1
}
for i := 1; i < N; i++ {
if rea2[i] != 0 {
if i <= n0 {
color[i] = 0
} else {
color[i] = 1
}
}
}
f(par2, V2, color)
centroidDecomposition2DFS(par0, V0, rea0, f)
centroidDecomposition2DFS(par1, V1, rea1, f)
}
func max(a, b int) int {
if a > b {
return a
}
return b
}
func min(a, b int) int {
if a < b {
return a
}
return b
}
// !Point Add Range Sum, 0-based.
type BITArray struct {
n int
total int
data []int
}
func NewBitArray(n int) *BITArray {
res := &BITArray{n: n, data: make([]int, n)}
return res
}
func NewBitArrayFrom(n int, f func(i int) int) *BITArray {
total := 0
data := make([]int, n)
for i := 0; i < n; i++ {
data[i] = f(i)
total += data[i]
}
for i := 1; i <= n; i++ {
j := i + (i & -i)
if j <= n {
data[j-1] += data[i-1]
}
}
return &BITArray{n: n, total: total, data: data}
}
func (b *BITArray) Add(index int, v int) {
b.total += v
for index++; index <= b.n; index += index & -index {
b.data[index-1] += v
}
}
// [0, end).
func (b *BITArray) QueryPrefix(end int) int {
if end > b.n {
end = b.n
}
res := 0
for ; end > 0; end -= end & -end {
res += b.data[end-1]
}
return res
}
// [start, end).
func (b *BITArray) QueryRange(start, end int) int {
if start < 0 {
start = 0
}
if end > b.n {
end = b.n
}
if start >= end {
return 0
}
if start == 0 {
return b.QueryPrefix(end)
}
pos, neg := 0, 0
for end > start {
pos += b.data[end-1]
end &= end - 1
}
for start > end {
neg += b.data[start-1]
start &= start - 1
}
return pos - neg
}
func (b *BITArray) QueryAll() int {
return b.total
}
func (b *BITArray) MaxRight(check func(index, preSum int) bool) int {
i := 0
s := 0
k := 1
for 2*k <= b.n {
k *= 2
}
for k > 0 {
if i+k-1 < b.n {
t := s + b.data[i+k-1]
if check(i+k, t) {
i += k
s = t
}
}
k >>= 1
}
return i
}
// 0/1 k(0-based) 1.
// UpperBound.
func (b *BITArray) Kth(k int) int {
return b.MaxRight(func(index, preSum int) bool { return preSum <= k })
}
func (b *BITArray) String() string {
sb := []string{}
for i := 0; i < b.n; i++ {
sb = append(sb, fmt.Sprintf("%d", b.QueryRange(i, i+1)))
}
return fmt.Sprintf("BitArray: [%v]", strings.Join(sb, ", "))
}
הההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההה
XXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXX
0