結果
問題 | No.778 クリスマスツリー |
ユーザー |
|
提出日時 | 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 |
ソースコード
package mainimport ("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) intdfs = 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/025644type DynamicSegmentTree[T any] struct {l, r introot *node[T]e func() Top 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.rootfor l, r := sg.l, sg.r; ; {if i == now.i {now.val = valbreak}if m := (l + r) / 2; i < m {if i > now.i {i, now.i = now.i, ival, now.val = now.val, val}p, now = now, now.lr = m} else {if i < now.i {i, now.i = now.i, ival, now.val = now.val, val}p, now = now, now.rl = m}if now == nil {now = newNode(i, val)now.p = pif now.i < p.i {p.l = now} else {p.r = now}now = pbreak}}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 intval, subVal Tp, 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 = 4096const maxBufSize = math.MaxInt64in := &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 = bwdefer func() {if err := bw.Flush(); err != nil {panic(fmt.Errorf("flush error: %w", err))}}()}outErr := &output{w: os.Stderr}loop := 1if 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 intfunc 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 := 1b := afor n > 0 {if n&1 > 0 {res *= b}n >>= 1b *= b}return res}func sum(s ...int) int {res := 0for _, v := range s {res += v}return res}// slice utility functionsfunc fillSlice[T any](s []T, v T) {for i := range s {s[i] = v}}func countSlice[T comparable](s []T, v T) int {res := 0for _, w := range s {if w == v {res++}}return res}