#948485 (Go) No.778 クリスマスツリー

提出ソース

結果

問題	No.778 クリスマスツリー
ユーザー	ynm3n
提出日時	2024-02-07 00:04:53
言語	Go (1.22.1)
結果	AC
実行時間	279 ms / 2,000 ms
コード長	7,306 bytes
コンパイル時間	13,279 ms
コンパイル使用メモリ	214,092 KB
実行使用メモリ	69,748 KB
最終ジャッジ日時	2024-02-07 00:05:12
合計ジャッジ時間	16,186 ms
ジャッジサーバーID （参考情報）	judge13 / judge15

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

テストケース

テストケース表示

入力	結果	実行時間実行使用メモリ
testcase_00	AC	1 ms 6,676 KB
testcase_01	AC	2 ms 6,676 KB
testcase_02	AC	2 ms 6,676 KB
testcase_03	AC	1 ms 6,676 KB
testcase_04	AC	1 ms 6,676 KB
testcase_05	AC	1 ms 6,676 KB
testcase_06	AC	187 ms 69,748 KB
testcase_07	AC	85 ms 18,612 KB
testcase_08	AC	279 ms 40,968 KB
testcase_09	AC	245 ms 20,564 KB
testcase_10	AC	254 ms 20,564 KB
testcase_11	AC	252 ms 20,560 KB
testcase_12	AC	239 ms 20,576 KB
testcase_13	AC	107 ms 19,300 KB
testcase_14	AC	202 ms 69,732 KB

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

ソースコード

raw source code

package main

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

// 解答欄
func solve(in *input, out, outErr *output) {
	n := in.nextInt()
	graph := make([][]int, n)
	for i := 1; i < n; i++ {
		a := in.nextInt()
		graph[a] = append(graph[a], i)
	}

	sg := NewDynamicSegmentTree[int](0, n, func() int { return 0 }, func(a, b int) int { return a + b })
	var dfs func(now int) int
	dfs = func(now int) int {
		res := sg.Product(0, now)
		sg.Set(now, 1)
		for _, next := range graph[now] {
			res += dfs(next)
		}
		sg.Set(now, 0)
		return res
	}

	ans := dfs(0)
	out.println(ans)
}

// ポインタを使う 必要な部分だけ作るやつ
// 参考にさせていただいた記事:
// https://kazuma8128.hatenablog.com/entry/2018/11/29/093827
// https://lorent-kyopro.hatenablog.com/entry/2021/03/12/025644
type DynamicSegmentTree[T any] struct {
	l, r int
	root *node[T]
	e    func() T
	op   func(a, b T) T
}

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

func NewDynamicSegmentTreeWith[T any](s []T, e func() T, op func(a, b T) T) *DynamicSegmentTree[T] {
	sg := NewDynamicSegmentTree(0, len(s), e, op)
	for i, val := range s {
		sg.Set(i, val)
	}
	return sg
}

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

	p := (*node[T])(nil)
	now := sg.root
	for l, r := sg.l, sg.r; ; {
		if i == now.i {
			now.val = val
			break
		}
		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 now == nil {
			now = newNode(i, val)
			now.p = p
			if now.i < p.i {
				p.l = now
			} else {
				p.r = now
			}
			now = p
			break
		}
	}

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

func (sg *DynamicSegmentTree[T]) Get(i int) T {
	sg.checkInRange(i)
	for now := sg.root; now != nil; {
		if i == now.i {
			return now.val
		} else if i < now.i {
			now = now.l
		} else {
			now = now.r
		}
	}
	return sg.e()
}

func (sg *DynamicSegmentTree[T]) Product(l, r int) T {
	sg.checkInRangeLR(l, r)
	if sg.root == nil || l == r {
		return sg.e()
	}
	return sg.product(sg.root, l, r, sg.l, sg.r)
}

func (sg *DynamicSegmentTree[T]) product(n *node[T], argL, argR, l, r int) T {
	if n == nil || r <= argL || argR <= l {
		return sg.e()
	}
	if argL <= l && r <= argR {
		return n.subVal
	}
	res := sg.product(n.l, argL, argR, l, n.i)
	if argL <= n.i && n.i < argR {
		res = sg.op(res, n.val)
	}
	res = sg.op(res, sg.product(n.r, argL, argR, n.i+1, r))
	return res
}

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

func (sg *DynamicSegmentTree[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))
}

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

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

func (sg *DynamicSegmentTree[T]) checkInRangeLR(argL, argR int) {
	if argL < sg.l || sg.r < argR {
		panic(fmt.Errorf("DynamicSegmentTree: index out of range: l=%d, r=%d, argL=%d, argR=%d", sg.l, sg.r, argL, argR))
	}
}

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

結果

テストケース

ソースコード