#934538 (Go) No.1038 TreeAddQuery

提出ソース

結果

問題	No.1038 TreeAddQuery
ユーザー	草苺奶昔
提出日時	2023-12-02 15:28:43
言語	Go (1.22.1)
結果	AC
実行時間	1,498 ms / 4,000 ms
コード長	11,707 bytes
コンパイル時間	13,517 ms
コンパイル使用メモリ	210,364 KB
実行使用メモリ	268,160 KB
最終ジャッジ日時	2023-12-02 15:29:30
合計ジャッジ時間	40,199 ms
ジャッジサーバーID （参考情報）	judge13 / judge9

このコードへのチャレンジ
（要ログイン）

テストケース

テストケース表示

入力	結果	実行時間実行使用メモリ
testcase_00	AC	1 ms 6,676 KB
testcase_01	AC	1 ms 6,676 KB
testcase_02	AC	1 ms 6,676 KB
testcase_03	AC	19 ms 8,468 KB
testcase_04	AC	20 ms 8,668 KB
testcase_05	AC	18 ms 8,700 KB
testcase_06	AC	15 ms 8,444 KB
testcase_07	AC	21 ms 8,560 KB
testcase_08	AC	1,369 ms 225,156 KB
testcase_09	AC	1,430 ms 244,784 KB
testcase_10	AC	1,416 ms 237,248 KB
testcase_11	AC	1,410 ms 255,780 KB
testcase_12	AC	1,391 ms 226,580 KB
testcase_13	AC	1,397 ms 206,996 KB
testcase_14	AC	1,388 ms 217,492 KB
testcase_15	AC	1,377 ms 211,804 KB
testcase_16	AC	1,426 ms 268,160 KB
testcase_17	AC	1,422 ms 229,812 KB
testcase_18	AC	789 ms 133,872 KB
testcase_19	AC	851 ms 145,372 KB
testcase_20	AC	854 ms 150,856 KB
testcase_21	AC	867 ms 138,236 KB
testcase_22	AC	949 ms 146,688 KB
testcase_23	AC	983 ms 190,284 KB
testcase_24	AC	1,152 ms 207,680 KB
testcase_25	AC	1,498 ms 234,688 KB
testcase_26	AC	1,174 ms 158,940 KB

権限があれば一括ダウンロードができます

ソースコード

raw source code

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 : 将root节点的值加上x (点权加)
// 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
// 给定一颗树，初始时点权为0，有q个操作.
// 每次操作输出顶点x的点权，并将距离x距离在[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, ", "))
}