次短路计数

main/1600-1699/1650G.go

@@ -0,0 +1,60 @@
+package main
+
+import (
+	. "fmt"
+	"io"
+)
+
+func cf1650G(in io.Reader, out io.Writer) {
+	const mod = 1_000_000_007
+	var T, n, m, s, t int
+	for Fscan(in, &T); T > 0; T-- {
+		Fscan(in, &n, &m, &s, &t)
+		s--
+		t--
+		g := make([][]int, n)
+		for ; m > 0; m-- {
+			var v, w int
+			Fscan(in, &v, &w)
+			v--
+			w--
+			g[v] = append(g[v], w)
+			g[w] = append(g[w], v)
+		}
+
+		// 偶数下标表示最短路，奇数下标表示次短路
+		f := make([]int, n*2)
+		dis := make([]int, n*2)
+		for i := range dis {
+			dis[i] = -1
+		}
+		dis[s*2] = 0
+		f[s*2] = 1
+		q := []int{s * 2}
+		for len(q) > 0 {
+			v := q[0]
+			q = q[1:]
+			newD := dis[v] + 1
+			for _, w := range g[v/2] {
+				w *= 2
+				if dis[w] < 0 {
+					dis[w] = newD
+					f[w] = f[v]
+					q = append(q, w)
+				} else if newD == dis[w] {
+					f[w] = (f[w] + f[v]) % mod
+				} else if newD == dis[w]+1 { // 次短路
+					w++
+					if dis[w] < 0 {
+						dis[w] = newD
+						q = append(q, w)
+					}
+					f[w] = (f[w] + f[v]) % mod
+				}
+			}
+		}
+		Fprintln(out, (f[t*2]+f[t*2+1])%mod)
+	}
+}
+
+//func main() { cf1650G(bufio.NewReader(os.Stdin), os.Stdout) }

main/1600-1699/1650G_test.go

@@ -0,0 +1,70 @@
+// Generated by copypasta/template/generator_test.go
+package main
+
+import (
+	"github.com/EndlessCheng/codeforces-go/main/testutil"
+	"testing"
+)
+
+// https://codeforces.com/problemset/problem/1650/G
+// https://codeforces.com/problemset/status/1650/problem/G
+func Test_cf1650G(t *testing.T) {
+	testCases := [][2]string{
+		{
+			`4
+
+4 4
+1 4
+1 2
+3 4
+2 3
+2 4
+
+6 8
+6 1
+1 4
+1 6
+1 5
+1 2
+5 6
+4 6
+6 3
+2 6
+
+5 6
+1 3
+3 5
+5 4
+3 1
+4 2
+2 1
+1 4
+
+8 18
+5 1
+2 1
+3 1
+4 2
+5 2
+6 5
+7 3
+8 4
+6 4
+8 7
+1 4
+4 7
+1 6
+6 7
+3 8
+8 5
+4 5
+4 3
+8 2`,
+			`2
+4
+1
+11`,
+		},
+	}
+	testutil.AssertEqualStringCase(t, testCases, 0, cf1650G)
+}