|
| 1 | +let m = 100000 |
| 2 | +var fact = [Int](repeating: 0, count: m + 1) |
| 3 | +var invFact = [Int](repeating: 0, count: m + 1) |
| 4 | +let mod = 1_000_000_007 |
| 5 | +var initValue = false |
| 6 | + |
| 7 | +class Solution { |
| 8 | + |
| 9 | + // Solution by Sergey Leschev |
| 10 | + // 2954. Count the Number of Infection Sequences |
| 11 | + |
| 12 | + func modPow(_ x: Int, _ y: Int, _ mod: Int) -> Int { |
| 13 | + if y == 0 { |
| 14 | + return 1 |
| 15 | + } |
| 16 | + var p = modPow(x, y / 2, mod) % mod |
| 17 | + p = (p * p) % mod |
| 18 | + return y % 2 == 1 ? (p * x) % mod : p |
| 19 | + } |
| 20 | + |
| 21 | + func binomCoeff(_ n: Int, _ k: Int) -> Int { |
| 22 | + return max(1, (fact[n] * invFact[k]) % mod * invFact[n - k] % mod) |
| 23 | + } |
| 24 | + |
| 25 | + func numberOfSequence(_ n: Int, _ sick: [Int]) -> Int { |
| 26 | + if !initValue { |
| 27 | + initValue = true |
| 28 | + fact[0] = 1 |
| 29 | + for i in 1...m { |
| 30 | + fact[i] = (fact[i - 1] * i) % mod |
| 31 | + } |
| 32 | + invFact[m] = modPow(fact[m], mod - 2, mod) |
| 33 | + for i in stride(from: m - 1, to: 0, by: -1) { |
| 34 | + invFact[i] = (invFact[i + 1] * (i + 1)) % mod |
| 35 | + } |
| 36 | + } |
| 37 | + |
| 38 | + var res: Int64 = 1 |
| 39 | + for i in 1..<sick.count { |
| 40 | + let group = sick[i] - sick[i - 1] - 1 |
| 41 | + res = (res * Int64(modPow(2, max(0, group - 1), mod))) % Int64(mod) // combine |
| 42 | + res = (res * Int64(binomCoeff(sick[i] - i, group))) % Int64(mod) // interweave |
| 43 | + } |
| 44 | + |
| 45 | + return Int(res * Int64(binomCoeff(n - sick.count, n - sick.last! - 1)) % Int64(mod)) |
| 46 | + } |
| 47 | +} |
0 commit comments