- Translate each rule into one explicit state update.
- Maintain the invariant after every processed item.
- Return the accumulated state once all relevant input has been handled.
Code notes
- 120 lines of Go from the credited upstream file 763A.go.
- The implementation visibly relies on sequence storage.
- No explicit loop blocks detected.
Complexity
Count the number and nesting of passes over the input, then include the maintained containers in the memory estimate.
Check the problem constraints before deciding whether this complexity will pass.
Use this to learn the idea, then write your own version.
1package main2 3import (4 "bufio"5 . "fmt"6 "io"7)8 910 1112func CF763A(_r io.Reader, _w io.Writer) {13 in := bufio.NewReader(_r)14 out := bufio.NewWriter(_w)15 defer out.Flush()16 17 var n, cnt int18 Fscan(in, &n)19 es := make([][2]int, n-1)20 for i := range es {21 Fscan(in, &es[i][0], &es[i][1])22 }23 c := make([]int, n)24 for i := range c {25 Fscan(in, &c[i])26 }27 diff := make([]int, n)28 for _, e := range es {29 if v, w := e[0]-1, e[1]-1; c[v] != c[w] {30 diff[v]++31 diff[w]++32 cnt++33 }34 }35 for i, d := range diff {36 if d == cnt {37 Fprint(out, "YES\n", i+1)38 return39 }40 }41 Fprint(out, "NO")42}43 4445func CF763Adp(_r io.Reader, _w io.Writer) {46 in := bufio.NewReader(_r)47 out := bufio.NewWriter(_w)48 defer out.Flush()49 50 var n, v, w int51 Fscan(in, &n)52 g := make([][]int, n)53 for i := 1; i < n; i++ {54 Fscan(in, &v, &w)55 v--56 w--57 g[v] = append(g[v], w)58 g[w] = append(g[w], v)59 }60 c := make([]int, n)61 for i := range c {62 Fscan(in, &c[i])63 }64 65 same := make([]bool, n)66 for i := range same {67 same[i] = true68 }69 var f func(v, fa int) bool70 f = func(v, fa int) bool {71 for _, w := range g[v] {72 if w != fa {73 if !f(w, v) || c[w] != c[v] {74 same[v] = false75 }76 }77 }78 return same[v]79 }80 f(0, -1)81 82 ans := -183 var f2 func(v, fa int)84 f2 = func(v, fa int) {85 cntDiff, allSame := 0, fa < 0 || c[v] == c[fa]86 for _, w := range g[v] {87 if w != fa {88 if !same[w] {89 cntDiff++90 } else if c[w] != c[v] {91 allSame = false92 }93 }94 }95 if cntDiff == 0 {96 ans = v97 return98 }99 if !allSame || cntDiff > 1 {100 return101 }102 for _, w := range g[v] {103 if w != fa {104 if !same[w] {105 f2(w, v)106 break107 }108 }109 }110 }111 f2(0, -1)112 if ans < 0 {113 Fprint(out, "NO")114 } else {115 Fprint(out, "YES\n", ans+1)116 }117}118 119120