#947693 (Go) No.789 範囲の合計

提出ソース
結果

問題	No.789 範囲の合計
コンテスト
ユーザー	ynm3n
提出日時	2024-02-03 17:57:03
言語	Go (1.23.4)
結果	AC
実行時間	90 ms / 1,000 ms
コード長	7,369 bytes
コンパイル時間	11,537 ms
コンパイル使用メモリ	226,040 KB
実行使用メモリ	6,944 KB
最終ジャッジ日時	2024-09-28 10:55:50
合計ジャッジ時間	13,468 ms
ジャッジサーバーID （参考情報）	judge4 / judge1
このコードへのチャレンジ
（要ログイン）
ファイルパターン	結果
other	AC * 15
権限があれば一括ダウンロードができます
ソースコード

raw source code
package main

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

// 解答欄
func solve(in *input, out, outErr *output) {
	n := in.nextInt()

	sg := NewSegmentTree(0, 1e9+1, func() int { return 0 }, func(a, b int) int { return a + b })
	ans := 0
	for i := 0; i < n; i++ {
		switch q := in.nextInt(); q {
		case 0:
			x, y := in.nextInt2()
			tmp := sg.Get(x)
			sg.Set(x, tmp+y)
		case 1:
			l, r := in.nextInt2()
			r++
			tmp := sg.Product(l, r)
			ans += tmp
		}
	}
	out.println(ans)
}

// ポインタを使う 必要な部分だけ作るやつ
type SegmentTree[T any] struct {
	l, r int
	root *node[T]

	e  func() T
	op func(a, b T) T
}

func NewSegmentTree[T any](l, r int, e func() T, op func(a, b T) T) *SegmentTree[T] {
	return &SegmentTree[T]{
		l:    l,
		r:    r,
		root: nil,
		e:    e,
		op:   op,
	}
}

func (sg *SegmentTree[T]) Set(i int, val T) {
	sg.checkInRange(i)
	if sg.root == nil {
		sg.root = newNode(i, val)
		return
	}

	var p, upd *node[T]
	now := sg.root
	for l, r := sg.l, sg.r; ; {
		if m := (l + r) / 2; i < m {
			if i > now.i {
				i, now.i = now.i, i
				val, now.val = now.val, val
			}
			p, now = now, now.l
			r = m
		} else {
			if i < now.i {
				i, now.i = now.i, i
				val, now.val = now.val, val
			}
			p, now = now, now.r
			l = m
		}
		if i == p.i {
			p.val = val
			upd = p
			break
		} else if now == nil {
			leaf := newNode(i, val)
			leaf.p = p
			if leaf.i < p.i {
				p.l = leaf
			} else {
				p.r = leaf
			}
			upd = p
			break
		}
	}

	for ; upd != nil; upd = upd.p {
		sg.update(upd)
	}
}

func (sg *SegmentTree[T]) Get(i int) T {
	sg.checkInRange(i)
	if n, _ := sg.search(i); n != nil && n.i == i {
		return n.val
	}
	return sg.e()
}

func (sg *SegmentTree[T]) Product(l, r int) T {
	sg.checkInRange(l)
	sg.checkInRange(r - 1)
	if sg.root == nil {
		return sg.e()
	}

	nL, depthL := sg.search(l)
	cL := nL.l
	nR, depthR := sg.search(r - 1)
	cR := nR.r

	valL, valR := sg.e(), sg.e()
	for nL != nR {
		if depthL >= depthR {
			// searchで辿る道を考えると、右の子から移動してきた場合、nLは必ず範囲外
			if cL == nL.l {
				if l <= nL.i && nL.i < r {
					valL = sg.op(valL, nL.val)
				}
				valL = sg.op(valL, sg.subtreeVal(nL.r))
			}
			nL, cL = nL.p, nL
			depthL--
		}
		if depthL < depthR {
			// searchで辿る道を考えると、左の子から移動してきた場合、nRは必ず範囲外
			if cR == nR.r {
				if l <= nR.i && nR.i < r {
					valR = sg.op(nR.val, valR)
				}
				valR = sg.op(sg.subtreeVal(nR.l), valR)
			}
			nR, cR = nR.p, nR
			depthR--
		}
	}
	if nL.i < l || r <= nL.i { // 範囲内のノードが存在しない場合はここで分岐
		return sg.e()
	}
	return sg.op(sg.op(valL, nL.val), valR)
}

func (sg *SegmentTree[T]) ProductAll() T {
	return sg.root.subVal
}

func (sg *SegmentTree[T]) update(n *node[T]) {
	n.subVal = sg.op(sg.subtreeVal(n.l), n.val)
	n.subVal = sg.op(n.subVal, sg.subtreeVal(n.r))
}

// インデックスがiに近いノードとそのノードの深さを返す
// ノードの個数が0の時、深さは-1になる
func (sg *SegmentTree[T]) search(i int) (*node[T], int) {
	p := (*node[T])(nil)
	depth := 0
	for now := sg.root; now != nil; {
		p = now
		if i == now.i {
			return now, depth
		} else if i < now.i {
			now = now.l
		} else {
			now = now.r
		}
		depth++
	}
	return p, depth - 1
}

func (sg *SegmentTree[T]) subtreeVal(n *node[T]) T {
	if n != nil {
		return n.subVal
	}
	return sg.e()
}

func (sg *SegmentTree[T]) checkInRange(i int) {
	if i < sg.l || sg.r <= i {
		panic(fmt.Sprintf("SegmentTree: index out of range: l=%d, r=%d, i=%d", sg.l, sg.r, i))
	}
}

type node[T any] struct {
	i           int
	val, subVal T
	p, l, r     *node[T]
}

func newNode[T any](i int, val T) *node[T] {
	return &node[T]{
		i:      i,
		val:    val,
		subVal: val,
	}
}

// 入出力の準備(in, out, outErr)
func main() {
	const (
		// インタラクティブ問題の時 true
		// バッファ無し出力に変更する。入力は変わらない
		interactiveProblem = false

		// 「X個のテストケースが与えられる」系の問題の時 true
		// 入力の一番最初でテストケース数が与えられる場合のみ使える
		variableSolveCount = false
	)

	const startBufSize = 4096
	const maxBufSize = math.MaxInt64
	in := &input{sc: bufio.NewScanner(os.Stdin)}
	in.sc.Buffer(make([]byte, startBufSize), maxBufSize)
	in.sc.Split(bufio.ScanWords)
	out := &output{}
	if interactiveProblem {
		out.w = os.Stdout
	} else {
		bw := bufio.NewWriter(os.Stdout)
		out.w = bw
		defer func() {
			if err := bw.Flush(); err != nil {
				panic(fmt.Errorf("flush error: %w", err))
			}
		}()
	}
	outErr := &output{w: os.Stderr}

	loop := 1
	if variableSolveCount {
		if _, err := fmt.Scan(&loop); err != nil {
			panic(fmt.Errorf("valiableSolveCount = %v, loop = %v: %w", variableSolveCount, loop, err))
		}
	}
	for i := 0; i < loop; i++ {
		solve(in, out, outErr)
	}
}

type input struct{ sc *bufio.Scanner }

func (in *input) nextString() string {
	in.sc.Scan()
	return in.sc.Text()
}
func (in *input) nextStrings(n int) []string {
	res := make([]string, n)
	for i := range res {
		res[i] = in.nextString()
	}
	return res
}
func (in *input) nextRunes() []rune {
	return []rune(in.nextString())
}
func (in *input) nextInt() int {
	s := in.nextString()
	res, err := strconv.Atoi(s)
	if err != nil {
		panic(fmt.Errorf("in.nextInt: %w", err))
	}
	return res
}
func (in *input) nextInts(n int) []int {
	res := make([]int, n)
	for i := range res {
		res[i] = in.nextInt()
	}
	return res
}
func (in *input) nextFloat() float64 {
	s := in.nextString()
	res, err := strconv.ParseFloat(s, 64)
	if err != nil {
		panic(fmt.Errorf("in.nextFloat: %w", err))
	}
	return res
}
func (in *input) nextFloats(n int) []float64 {
	res := make([]float64, n)
	for i := range res {
		res[i] = in.nextFloat()
	}
	return res
}
func (in *input) nextInt2() (int, int) {
	return in.nextInt(), in.nextInt()
}
func (in *input) nextInt3() (int, int, int) {
	return in.nextInt(), in.nextInt(), in.nextInt()
}
func (in *input) nextInt4() (int, int, int, int) {
	return in.nextInt(), in.nextInt(), in.nextInt(), in.nextInt()
}

type output struct{ w io.Writer }

func (out *output) print(a ...any) {
	if _, err := fmt.Fprint(out.w, a...); err != nil {
		panic(fmt.Errorf("out.print: %w", err))
	}
}
func (out *output) printf(format string, a ...any) { out.print(fmt.Sprintf(format, a...)) }
func (out *output) println(a ...any)               { out.print(fmt.Sprintln(a...)) }

// simple math functions for int
func max(as ...int) int {
	res := as[0]
	for _, a := range as {
		if res < a {
			res = a
		}
	}
	return res
}
func min(as ...int) int {
	res := as[0]
	for _, a := range as {
		if res > a {
			res = a
		}
	}
	return res
}
func chMax(a *int, b int) {
	*a = max(*a, b)
}
func chMin(a *int, b int) {
	*a = min(*a, b)
}
func abs(a int) int {
	if a < 0 {
		return -a
	}
	return a
}
func pow(a, n int) int {
	res := 1
	b := a
	for n > 0 {
		if n&1 > 0 {
			res *= b
		}
		n >>= 1
		b *= b
	}
	return res
}
func sum(s ...int) int {
	res := 0
	for _, v := range s {
		res += v
	}
	return res
}

// slice utility functions
func fillSlice[T any](s []T, v T) {
	for i := range s {
		s[i] = v
	}
}
func countSlice[T comparable](s []T, v T) int {
	res := 0
	for _, w := range s {
		if w == v {
			res++
		}
	}
	return res
}
yukicoder

結果

ソースコード
