CF1923E 树形DP

main/1900-1999/1923E.go

@@ -0,0 +1,52 @@
+package main
+
+import (
+	"bufio"
+	. "fmt"
+	"io"
+)
+
+// https://space.bilibili.com/206214
+func cf1923E(_r io.Reader, _w io.Writer) {
+	in := bufio.NewReader(_r)
+	out := bufio.NewWriter(_w)
+	defer out.Flush()
+
+	var T, n, v, w int
+	for Fscan(in, &T); T > 0; T-- {
+		Fscan(in, &n)
+		cs := make([]int, n)
+		for i := range cs {
+			Fscan(in, &cs[i])
+		}
+		g := make([][]int, n)
+		for i := 1; i < n; i++ {
+			Fscan(in, &v, &w)
+			v--
+			w--
+			g[v] = append(g[v], w)
+			g[w] = append(g[w], v)
+		}
+
+		ans := 0
+		cnt := make([]int, n+1)
+		var dfs func(int, int)
+		dfs = func(v, fa int) {
+			c := cs[v]
+			tmp := cnt[c]
+			ans += tmp
+			for _, w := range g[v] {
+				if w != fa {
+					cnt[c] = 1
+					dfs(w, v)
+				}
+			}
+			cnt[c] = tmp + 1
+			return
+		}
+		dfs(0, -1)
+		Fprintln(out, ans)
+	}
+}
+
+//func main() { cf1923E(os.Stdin, os.Stdout) }

main/1900-1999/1923E_test.go

@@ -0,0 +1,43 @@
+// Generated by copypasta/template/generator_test.go
+package main
+
+import (
+	"github.com/EndlessCheng/codeforces-go/main/testutil"
+	"testing"
+)
+
+// https://codeforces.com/contest/1923/problem/E
+// https://codeforces.com/problemset/status/1923/problem/E
+func Test_cf1923E(t *testing.T) {
+	testCases := [][2]string{
+		{
+			`4
+3
+1 2 1
+1 2
+2 3
+5
+2 1 2 1 2
+1 2
+1 3
+3 4
+4 5
+5
+1 2 3 4 5
+1 2
+1 3
+3 4
+4 5
+4
+2 2 2 2
+3 1
+3 2
+3 4`,
+			`1
+3
+0
+3`,
+		},
+	}
+	testutil.AssertEqualStringCase(t, testCases, 0, cf1923E)
+}