結果

問題 No.778 クリスマスツリー
ユーザー ynm3n
提出日時 2024-02-07 00:04:53
言語 Go
(1.23.4)
結果
AC  
実行時間 390 ms / 2,000 ms
コード長 7,306 bytes
コンパイル時間 15,866 ms
コンパイル使用メモリ 232,180 KB
実行使用メモリ 64,896 KB
最終ジャッジ日時 2024-09-28 12:19:14
合計ジャッジ時間 17,799 ms
ジャッジサーバーID
(参考情報)
judge1 / judge4
このコードへのチャレンジ
(要ログイン)
ファイルパターン 結果
sample AC * 3
other AC * 12
権限があれば一括ダウンロードができます

ソースコード

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

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
}
הההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההה
XXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXX
0