#878321 (Go) No.1340 おーじ君をさがせ

提出ソース

結果

問題	No.1340 おーじ君をさがせ
ユーザー	草苺奶昔
提出日時	2023-06-05 19:53:08
言語	Go (1.22.1)
結果	AC
実行時間	21 ms / 2,000 ms
コード長	9,267 bytes
コンパイル時間	11,269 ms
コンパイル使用メモリ	211,840 KB
実行使用メモリ	8,072 KB
最終ジャッジ日時	2023-08-28 09:41:39
合計ジャッジ時間	13,493 ms
ジャッジサーバーID （参考情報）	judge13 / judge12

このコードへのチャレンジ
（要ログイン）

テストケース

テストケース表示

入力	結果	実行時間実行使用メモリ
testcase_00	AC	2 ms 4,352 KB
testcase_01	AC	1 ms 4,356 KB
testcase_02	AC	1 ms 4,356 KB
testcase_03	AC	1 ms 4,352 KB
testcase_04	AC	1 ms 4,356 KB
testcase_05	AC	1 ms 4,356 KB
testcase_06	AC	1 ms 4,360 KB
testcase_07	AC	1 ms 4,356 KB
testcase_08	AC	2 ms 4,352 KB
testcase_09	AC	1 ms 4,356 KB
testcase_10	AC	2 ms 4,352 KB
testcase_11	AC	6 ms 5,680 KB
testcase_12	AC	3 ms 4,356 KB
testcase_13	AC	5 ms 5,668 KB
testcase_14	AC	6 ms 5,692 KB
testcase_15	AC	6 ms 5,688 KB
testcase_16	AC	3 ms 4,356 KB
testcase_17	AC	3 ms 4,356 KB
testcase_18	AC	5 ms 5,676 KB
testcase_19	AC	2 ms 4,356 KB
testcase_20	AC	2 ms 4,360 KB
testcase_21	AC	13 ms 5,676 KB
testcase_22	AC	14 ms 7,812 KB
testcase_23	AC	13 ms 5,672 KB
testcase_24	AC	16 ms 8,060 KB
testcase_25	AC	10 ms 4,352 KB
testcase_26	AC	7 ms 4,352 KB
testcase_27	AC	6 ms 4,356 KB
testcase_28	AC	9 ms 4,352 KB
testcase_29	AC	5 ms 4,352 KB
testcase_30	AC	13 ms 5,672 KB
testcase_31	AC	20 ms 8,068 KB
testcase_32	AC	20 ms 8,068 KB
testcase_33	AC	20 ms 8,064 KB
testcase_34	AC	18 ms 8,072 KB
testcase_35	AC	21 ms 8,072 KB
testcase_36	AC	2 ms 4,352 KB
testcase_37	AC	8 ms 4,352 KB
testcase_38	AC	20 ms 8,064 KB
testcase_39	AC	11 ms 4,352 KB
testcase_40	AC	11 ms 4,352 KB
testcase_41	AC	11 ms 4,356 KB
testcase_42	AC	2 ms 4,352 KB
testcase_43	AC	1 ms 4,356 KB
testcase_44	AC	1 ms 4,356 KB
testcase_45	AC	1 ms 4,352 KB
testcase_46	AC	5 ms 4,356 KB
testcase_47	AC	5 ms 4,352 KB
testcase_48	AC	5 ms 4,352 KB
testcase_49	AC	4 ms 4,356 KB
testcase_50	AC	4 ms 4,356 KB
testcase_51	AC	4 ms 4,352 KB
testcase_52	AC	6 ms 5,688 KB
testcase_53	AC	6 ms 5,684 KB
testcase_54	AC	6 ms 5,688 KB
testcase_55	AC	6 ms 5,672 KB
testcase_56	AC	1 ms 4,356 KB
testcase_57	AC	1 ms 4,356 KB
testcase_58	AC	1 ms 4,356 KB
testcase_59	AC	3 ms 4,356 KB
testcase_60	AC	2 ms 4,356 KB
testcase_61	AC	3 ms 4,352 KB

権限があれば一括ダウンロードができます

ソースコード

raw source code

package main

import (
	"bufio"
	"fmt"
	"math/bits"
	"math/rand"
	"os"
	"strings"
	"time"
)

func main() {
	yuki1340()
	// test()
	// bs := NewBitset(70)
	// bs.Set(10)
	// bs.Set(15)
	// fmt.Println(bs.Has(10))
	// fmt.Println(bs._RangeHas(10, 16))
	// bs.Set(63)
	// bs.Set(64)
	// bs.Set(65)
	// fmt.Println(bs._RangeHas(62, 66))
}

// https://yukicoder.me/problems/no/1340
// 给定一个n个点m条边的有向图，求t步后可能所在的顶点个数(每一步必须移动到一个相邻点).
// n<=100 m<=1e4 t<=1e18
func yuki1340() {

	in := bufio.NewReader(os.Stdin)
	out := bufio.NewWriter(os.Stdout)
	defer out.Flush()

	var n, m, t int
	fmt.Fscan(in, &n, &m, &t)
	mat := NewBooleanSquareMatrix(n)
	for i := 0; i < m; i++ {
		var a, b int
		fmt.Fscan(in, &a, &b)
		mat.Set(a, b, true)
	}
	mat.IPow(t)

	res := 0
	for i := 0; i < n; i++ {
		if mat.Get(0, i) {
			res++
		}
	}
	fmt.Fprintln(out, res)
}

// https://leetcode.cn/problems/course-schedule-iv/
func checkIfPrerequisite(numCourses int, prerequisites [][]int, queries [][]int) []bool {
	mat := NewBooleanSquareMatrix(numCourses)
	for _, p := range prerequisites {
		mat.Set(p[0], p[1], true)
	}
	trans := mat.TransitiveClosure()
	res := make([]bool, len(queries))
	for i, q := range queries {
		res[i] = trans.Get(q[0], q[1])
	}
	return res
}

func test() {
	// ====================
	// 测试随机矩阵
	// 5000*5000的矩阵乘法:293.2131ms
	// 2000*2000的传递闭包:398.609ms
	// ====================
	// 测试稀疏矩阵
	// 5000*5000的矩阵乘法:247.3379ms
	// 2000*2000的传递闭包:394.9004ms
	// ====================
	// 测试稠密矩阵
	// 5000*5000的矩阵乘法:246.341ms
	// 2000*2000的传递闭包:395.3344ms

	mat := NewBooleanSquareMatrix(3)
	mat.Set(0, 0, true)
	mat.Set(0, 1, true)
	mat.Set(1, 2, true)
	mat.Set(1, 0, true)

	testRandom := func() {
		fmt.Println(strings.Repeat("=", 20))
		fmt.Println("测试随机矩阵")
		// !随机01矩阵
		// 5000*5000的矩阵乘法 => 697.955ms
		N_5000 := 5000
		mat := NewBooleanSquareMatrix(N_5000)
		for i := 0; i < N_5000; i++ {
			for j := 0; j < N_5000; j++ {
				if rand.Intn(2) == 0 {
					mat.Set(i, j, true)
				}
			}
		}
		time1 := time.Now()
		Mul(mat, mat)
		time2 := time.Now()
		fmt.Println(fmt.Sprintf("5000*5000的矩阵乘法:%v", time2.Sub(time1)))

		// 2000*2000的传递闭包 => 830.0099ms
		N_2000 := 2000
		mat = NewBooleanSquareMatrix(N_2000)
		for i := 0; i < N_2000; i++ {
			for j := 0; j < N_2000; j++ {
				if rand.Intn(2) == 0 {
					mat.Set(i, j, true)
				}
			}
		}
		time3 := time.Now()
		mat.TransitiveClosure()
		time4 := time.Now()
		fmt.Println(fmt.Sprintf("2000*2000的传递闭包:%v", time4.Sub(time3)))
	}

	testSparse := func() {
		fmt.Println(strings.Repeat("=", 20))
		fmt.Println("测试稀疏矩阵")
		// !稀疏矩阵
		// 5000*5000的矩阵乘法 => 548.3657ms
		N_5000 := 5000
		mat := NewBooleanSquareMatrix(N_5000)
		for i := 0; i < N_5000; i++ {
			for j := 0; j < N_5000; j++ {
				if rand.Intn(10) == 0 {
					mat.Set(i, j, true)
				}
			}
		}
		time1 := time.Now()
		Mul(mat, mat)
		time2 := time.Now()
		fmt.Println(fmt.Sprintf("5000*5000的矩阵乘法:%v", time2.Sub(time1)))

		// 2000*2000的传递闭包 => 817.5958ms
		N_2000 := 2000
		mat = NewBooleanSquareMatrix(N_2000)
		for i := 0; i < N_2000; i++ {
			for j := 0; j < N_2000; j++ {
				if rand.Intn(10) == 0 {
					mat.Set(i, j, true)
				}
			}
		}
		time3 := time.Now()
		mat.TransitiveClosure()
		time4 := time.Now()
		fmt.Println(fmt.Sprintf("2000*2000的传递闭包:%v", time4.Sub(time3)))
	}

	testDense := func() {
		fmt.Println(strings.Repeat("=", 20))
		fmt.Println("测试稠密矩阵")
		// !稠密矩阵
		// 5000*5000的矩阵乘法 => 491.0388ms
		N_5000 := 5000
		mat := NewBooleanSquareMatrix(N_5000)
		for i := 0; i < N_5000; i++ {
			for j := 0; j < N_5000; j++ {
				mat.Set(i, j, true)
			}
		}
		time1 := time.Now()
		Mul(mat, mat)
		time2 := time.Now()
		fmt.Println(fmt.Sprintf("5000*5000的矩阵乘法:%v", time2.Sub(time1)))

		// 2000*2000的传递闭包 => 827.1923ms
		N_2000 := 2000
		mat = NewBooleanSquareMatrix(N_2000)
		for i := 0; i < N_2000; i++ {
			for j := 0; j < N_2000; j++ {
				mat.Set(i, j, true)
			}
		}
		time3 := time.Now()
		mat.TransitiveClosure()
		time4 := time.Now()
		fmt.Println(fmt.Sprintf("2000*2000的传递闭包:%v", time4.Sub(time3)))
	}

	testRandom()
	testSparse()
	testDense()
}

// trailing zero table
var _BSF [1e4 + 10]int

func init() {
	for i := range _BSF {
		_BSF[i] = bits.TrailingZeros(uint(i))
	}
}

// 布尔方阵.
type BooleanSquareMatrix struct {
	N  int
	bs []BitSet64
	dp []BitSet64 // 在计算矩阵乘法时用到
}

// n<=1e4.
func NewBooleanSquareMatrix(n int) *BooleanSquareMatrix {
	bs := make([]BitSet64, n)
	for i := range bs {
		bs[i] = NewBitset(n)
	}
	return &BooleanSquareMatrix{N: n, bs: bs}
}

// n<=1e4.
func Eye(n int) *BooleanSquareMatrix {
	res := NewBooleanSquareMatrix(n)
	for i := 0; i < n; i++ {
		res.bs[i].Set(i)
	}
	return res
}

func Pow(mat *BooleanSquareMatrix, k int) *BooleanSquareMatrix {
	return mat.Copy().IPow(k)
}

func Mul(mat1, mat2 *BooleanSquareMatrix) *BooleanSquareMatrix {
	return mat1.Copy().IMul(mat2)
}

func Add(mat1, mat2 *BooleanSquareMatrix) *BooleanSquareMatrix {
	return mat1.Copy().IAdd(mat2)
}

// (A + I)^n 是传递闭包.
func (bm *BooleanSquareMatrix) TransitiveClosure() *BooleanSquareMatrix {
	n := bm.N
	newMat := Eye(n).IAdd(bm)
	newMat.IPow(n)
	return newMat
}

func (bm *BooleanSquareMatrix) IPow(k int) *BooleanSquareMatrix {
	res := Eye(bm.N)
	for k > 0 {
		if k&1 == 1 {
			res.IMul(bm)
		}
		bm.IMul(bm)
		k >>= 1
	}
	res.bs, bm.bs = bm.bs, res.bs
	return bm
}

// O(n^3/wlogn),这里logn指的是分块的大小.
func (bm *BooleanSquareMatrix) IMul(mat *BooleanSquareMatrix) *BooleanSquareMatrix {
	n := mat.N
	res := NewBooleanSquareMatrix(n)
	step := 8 // !理论最优是logn,实际取8效果最好(n为5000时)
	bm._initDpIfAbsent(step, n)
	dp := bm.dp

	for l, r := 0, step; l != n; l, r = r, r+step {
		if r > n {
			r = n
		}

		for s := 1; s < (1 << step); s++ {
			bsf := _BSF[s]
			if l+bsf < n {
				dp[s] = Or(dp[s^(1<<bsf)], mat.bs[l+bsf]) // Xor => f2矩阵乘法
			} else {
				dp[s] = dp[s^(1<<bsf)]
			}
		}

		for i := 0; i != n; i++ {
			// !这里是瓶颈
			// TODO:位运算优化
			now := int(bm.bs[i]._HasRange(l, r))
			res.bs[i].IOr(dp[now]) // IXor => f2矩阵乘法
		}
	}

	bm.bs, res.bs = res.bs, bm.bs
	return res
}

func (bm *BooleanSquareMatrix) IAdd(mat *BooleanSquareMatrix) *BooleanSquareMatrix {
	for i := 0; i < bm.N; i++ {
		bm.bs[i].IOr(mat.bs[i])
	}
	return bm
}

func (bm *BooleanSquareMatrix) Copy() *BooleanSquareMatrix {
	bs := make([]BitSet64, bm.N)
	for i := range bs {
		bs[i] = bm.bs[i].Copy()
	}
	return &BooleanSquareMatrix{N: bm.N, bs: bs, dp: bm.dp}
}

func (bm *BooleanSquareMatrix) Get(row, col int) bool {
	return bm.bs[row].Has(col)
}

func (bm *BooleanSquareMatrix) Set(row, col int, b bool) {
	if b {
		bm.bs[row].Set(col)
	} else {
		bm.bs[row].Reset(col)
	}
}

// To 2D grid.
func (mat *BooleanSquareMatrix) String() string {
	n := mat.N
	grid := make([][]int, n)
	for i := 0; i < n; i++ {
		grid[i] = make([]int, n)
		for j := 0; j < n; j++ {
			if mat.Get(i, j) {
				grid[i][j] = 1
			} else {
				grid[i][j] = 0
			}
		}
	}

	sb := strings.Builder{}
	sb.WriteString(fmt.Sprintf("BooleanSquareMatrix(%d,%d)\n", n, n))
	for i := 0; i < n; i++ {
		for j := 0; j < n; j++ {
			sb.WriteString(fmt.Sprintf("%d ", grid[i][j]))
		}
		sb.WriteString("\n")
	}

	return sb.String()
}

func (mat *BooleanSquareMatrix) _initDpIfAbsent(step int, n int) {
	if mat.dp == nil {
		dp := make([]BitSet64, 1<<step)
		for i := range dp {
			dp[i] = NewBitset(n)
		}
		mat.dp = dp
	}
}

type BitSet64 []uint64

func NewBitset(n int) BitSet64 { return make(BitSet64, n>>6+1) } // (n+_w-1)>>6

func (b BitSet64) Has(p int) bool { return b[p>>6]&(1<<(p&63)) != 0 }
func (b BitSet64) Flip(p int)     { b[p>>6] ^= 1 << (p & 63) }
func (b BitSet64) Set(p int)      { b[p>>6] |= 1 << (p & 63) }
func (b BitSet64) Reset(p int)    { b[p>>6] &^= 1 << (p & 63) }

func (b BitSet64) Copy() BitSet64 {
	res := make(BitSet64, len(b))
	copy(res, b)
	return res
}

func (b BitSet64) BitCount() int {
	res := 0
	for _, v := range b {
		res += bits.OnesCount64(v)
	}
	return res
}

// 将 c 的元素合并进 b
func (b BitSet64) IOr(c BitSet64) BitSet64 {
	for i, v := range c {
		b[i] |= v
	}
	return b
}

// !f2上的加法
func (b BitSet64) IXOr(c BitSet64) {
	for i, v := range c {
		b[i] ^= v
	}
}

func Or(a, b BitSet64) BitSet64 {
	res := make(BitSet64, len(a))
	for i, v := range a {
		res[i] = v | b[i]
	}
	return res
}

func Xor(a, b BitSet64) BitSet64 {
	res := make(BitSet64, len(a))
	for i, v := range a {
		res[i] = v ^ b[i]
	}
	return res
}

// ![l,r) 范围内数与bitset相交的数.r-l<64.
// eg: l=10, r=16
//    [15,14,13,12,11,10] & Bitset(10,11,13,15) => 101011
func (b BitSet64) _HasRange(l, r int) uint64 {
	posL, shiftL := l>>6, l&63
	posR, shiftR := r>>6, r&63
	maskL, maskR := ^(^uint64(0) << shiftL), ^(^uint64(0) << shiftR) // 低位全1
	if posL == posR {
		return (b[posL] & maskR) >> shiftL
	}
	// divL+1 == divR
	return (b[posL] & ^maskL)>>shiftL | (b[posR]&maskR)<<(63-shiftL)
}