CF466C

main/400-499/466C.go

@@ -6,51 +6,30 @@ import (
 	"io"
 )
 
-func Sol466C(reader io.Reader, writer io.Writer) {
-	in := bufio.NewReader(reader)
-	out := bufio.NewWriter(writer)
-	defer out.Flush()
-
-	var n int
+// https://space.bilibili.com/206214
+func cf466C(_r io.Reader, out io.Writer) {
+	in := bufio.NewReader(_r)
+	var n, cnt, ans int
 	Fscan(in, &n)
-	if n < 3 {
-		Fprintln(out, 0)
-		return
-	}
-	sum := make([]int64, n+1)
+	s := make([]int, n+1)
 	for i := 1; i <= n; i++ {
-		var val int64
-		Fscan(in, &val)
-		sum[i] = sum[i-1] + val
-	}
-	allSum := sum[n]
-
-	if allSum%3 != 0 {
-		Fprintln(out, 0)
-		return
+		Fscan(in, &s[i])
+		s[i] += s[i-1]
 	}
-	if allSum == 0 {
-		cnt0 := int64(0)
-		for _, s := range sum[1:n] {
-			if s == 0 {
-				cnt0++
-			}
-		}
-		Fprintln(out, cnt0*(cnt0-1)/2)
+	tot := s[n]
+	if tot%3 != 0 {
+		Fprint(out, 0)
 		return
 	}
-	partSum := allSum / 3
-	var ans, cnt1 int64
 	for i := 1; i < n; i++ {
-		if sum[i] == partSum {
-			cnt1++
-		} else if sum[i] == 2*partSum {
-			ans += cnt1
+		if s[i] == tot/3*2 {
+			ans += cnt
+		}
+		if s[i] == tot/3 {
+			cnt++
 		}
 	}
-	Fprintln(out, ans)
+	Fprint(out, ans)
 }
 
-//func main() {
-//	Sol466C(os.Stdin, os.Stdout)
-//}
+//func main() { cf466C(os.Stdin, os.Stdout) }

main/400-499/466C_test.go

@@ -1,31 +1,30 @@
+// Generated by copypasta/template/generator_test.go
 package main
 
 import (
 	"github.com/EndlessCheng/codeforces-go/main/testutil"
 	"testing"
 )
 
-func TestSol466C(t *testing.T) {
-	// just copy from website
-	rawText := `
-4
-0 0 0 0
-outputCopy
-3
-inputCopy
-5
-1 2 3 0 3
-outputCopy
-2
-inputCopy
-4
-0 1 -1 0
-outputCopy
-1
-inputCopy
-2
-4 1
-outputCopy
-0`
-	testutil.AssertEqual(t, rawText, Sol466C)
+// https://codeforces.com/problemset/problem/466/C
+// https://codeforces.com/problemset/status/466/problem/C
+func Test_cf466C(t *testing.T) {
+	testCases := [][2]string{
+		{
+			`5
+1 2 3 0 3`,
+			`2`,
+		},
+		{
+			`4
+0 1 -1 0`,
+			`1`,
+		},
+		{
+			`2
+4 1`,
+			`0`,
+		},
+	}
+	testutil.AssertEqualStringCase(t, testCases, 0, cf466C)
 }