CF1139C

main/1100-1199/1139C.go

@@ -0,0 +1,58 @@
+package main
+
+import (
+	"bufio"
+	. "fmt"
+	"io"
+)
+
+// https://space.bilibili.com/206214
+func CF1139C(_r io.Reader, out io.Writer) {
+	in := bufio.NewReader(_r)
+	const mod = 1_000_000_007
+	pow := func(x int64, n int) (res int64) {
+		x %= mod
+		res = 1
+		for ; n > 0; n /= 2 {
+			if n%2 > 0 {
+				res = res * x % mod
+			}
+			x = x * x % mod
+		}
+		return
+	}
+
+	var n, k, v, w, c int
+	Fscan(in, &n, &k)
+	g := make([][]int, n)
+	for i := 1; i < n; i++ {
+		Fscan(in, &v, &w, &c)
+		if c == 0 {
+			v--
+			w--
+			g[v] = append(g[v], w)
+			g[w] = append(g[w], v)
+		}
+	}
+	vis := make([]bool, n)
+	var f func(int) int
+	f = func(v int) int {
+		vis[v] = true
+		sz := 1
+		for _, w := range g[v] {
+			if !vis[w] {
+				sz += f(w)
+			}
+		}
+		return sz
+	}
+	ans := pow(int64(n), k)
+	for i, b := range vis {
+		if !b {
+			ans -= pow(int64(f(i)), k)
+		}
+	}
+	Fprint(out, (ans%mod+mod)%mod)
+}
+
+//func main() { CF1139C(os.Stdin, os.Stdout) }

main/1100-1199/1139C_test.go

@@ -0,0 +1,34 @@
+package main
+
+import (
+	"github.com/EndlessCheng/codeforces-go/main/testutil"
+	"testing"
+)
+
+// https://codeforces.com/problemset/problem/1139/C
+// https://codeforces.com/problemset/status/1139/problem/C
+func TestCF1139C(t *testing.T) {
+	// just copy from website
+	rawText := `
+inputCopy
+4 4
+1 2 1
+2 3 1
+3 4 1
+outputCopy
+252
+inputCopy
+4 6
+1 2 0
+1 3 0
+1 4 0
+outputCopy
+0
+inputCopy
+3 5
+1 2 1
+2 3 0
+outputCopy
+210`
+	testutil.AssertEqualCase(t, rawText, 0, CF1139C)
+}