← Home
For problem statement at 1000-1999/1100-1199/1170-1179/1172/problemC2.txt this is a correct solution, but verifier at 1000-1999/1100-1199/1170-1179/1172/verifierC2.go ends with REFERENCE_SOURCE_PATH not set
exit status 1 can you fix the verifier? package main

import (
	"bufio"
	"fmt"
	"os"
)

const MOD = 998244353

func power(base, exp int64) int64 {
	res := int64(1)
	base %= MOD
	for exp > 0 {
		if exp%2 == 1 {
			res = res * base % MOD
		}
		base = base * base % MOD
		exp /= 2
	}
	return res
}

func main() {
	reader := bufio.NewReader(os.Stdin)
	var n, m int
	fmt.Fscan(reader, &n, &m)

	a := make([]int, n)
	for i := 0; i < n; i++ {
		fmt.Fscan(reader, &a[i])
	}

	w := make([]int64, n)
	var W1, W0 int64
	for i := 0; i < n; i++ {
		fmt.Fscan(reader, &w[i])
		if a[i] == 1 {
			W1 += w[i]
		} else {
			W0 += w[i]
		}
	}

	invArray := make([]int64, 2*m+1)
	for idx := 0; idx <= 2*m; idx++ {
		d := W1 + W0 - int64(m) + int64(idx)
		if d > 0 {
			invArray[idx] = power(d, MOD-2)
		}
	}

	dp := make([]int64, m+1)
	nextDp := make([]int64, m+1)
	dp[0] = 1

	for i := 0; i < m; i++ {
		for j := 0; j <= i+1; j++ {
			nextDp[j] = 0
		}
		for j := 0; j <= i; j++ {
			if dp[j] == 0 {
				continue
			}
			invDen := invArray[2*j-i+m]

			prob1 := (W1 + int64(j)) % MOD * invDen % MOD
			nextDp[j+1] = (nextDp[j+1] + dp[j]*prob1) % MOD

			w0_current := W0 - (int64(i) - int64(j))
			if w0_current > 0 {
				prob0 := (w0_current % MOD) * invDen % MOD
				nextDp[j] = (nextDp[j] + dp[j]*prob0) % MOD
			}
		}
		dp, nextDp = nextDp, dp
	}

	var eW1, eW0 int64
	for j := 0; j <= m; j++ {
		if dp[j] > 0 {
			val1 := (W1 + int64(j)) % MOD
			eW1 = (eW1 + dp[j]*val1) % MOD

			w0_end := W0 - (int64(m) - int64(j))
			if w0_end > 0 {
				val0 := w0_end % MOD
				eW0 = (eW0 + dp[j]*val0) % MOD
			}
		}
	}

	ans := make([]int64, n)
	invW1 := power(W1, MOD-2)
	var invW0 int64
	if W0 > 0 {
		invW0 = power(W0, MOD-2)
	}

	for k := 0; k < n; k++ {
		if a[k] == 1 {
			ans[k] = w[k] % MOD * eW1 % MOD * invW1 % MOD
		} else {
			if W0 > 0 {
				ans[k] = w[k] % MOD * eW0 % MOD * invW0 % MOD
			} else {
				ans[k] = 0
			}
		}
	}

	writer := bufio.NewWriter(os.Stdout)
	for i := 0; i < n; i++ {
		fmt.Fprintf(writer, "%d ", ans[i])
	}
	fmt.Fprintln(writer)
	writer.Flush()
}