Advertisement
Not a member of Pastebin yet?
Sign Up,
it unlocks many cool features!
- MOD = 10**9 + 7
- VOWELS = "aeiou"
- ALLOWED_NEXT = {
- "a": "e",
- "e": "ai",
- "i": "aeou",
- "o": "iu",
- "u": "a",
- }
- ALLOWED_NEXT_IDXES = {
- VOWELS.find(k): [VOWELS.find(v) for v in vs]
- for k, vs in ALLOWED_NEXT.items()
- }
- def gen_update_matrix() -> List[List[int]]:
- ans = [[0] * len(VOWELS) for _ in range(len(VOWELS))]
- for k, vs in ALLOWED_NEXT_IDXES.items():
- for v in vs:
- ans[v][k] = 1
- return ans
- def matrix_mult(a: List[List[int]], b: List[List[int]]) -> List[List[int]]:
- a_rows = len(a)
- a_cols = len(a[0])
- b_rows = len(b)
- b_cols = len(b[0])
- assert a_cols == b_rows, f"{a_rows}x{a_cols} {b_rows}x{b_cols}"
- ans = [[0] * b_cols for _ in range(a_rows)]
- for cr in range(a_rows):
- for cb in range(b_cols):
- for i in range(a_cols):
- ans[cr][cb] += a[cr][i] * b[i][cb]
- ans[cr][cb] %= MOD
- return ans
- def matrix_pow_slow(m: List[List[int]], p: int) -> List[List[int]]:
- ans = m
- for i in range(p - 1):
- ans = matrix_mult(ans, m)
- return ans
- def matrix_pow(m: List[List[int]], p: int) -> List[List[int]]:
- assert len(m) == len(m[0])
- ans = [[0] * len(m) for _ in range(len(m))]
- for i in range(len(m)):
- ans[i][i] = 1
- cur_mult = m
- cur_pow = 1
- while cur_pow <= p:
- if p & cur_pow:
- ans = matrix_mult(ans, cur_mult)
- cur_pow *= 2
- cur_mult = matrix_mult(cur_mult, cur_mult)
- return ans
- class Solution:
- def countVowelPermutation(self, n: int) -> int:
- if n == 1:
- return len(VOWELS)
- initial = [[1] * len(VOWELS)]
- m = matrix_pow(gen_update_matrix(), n - 1)
- return sum(matrix_mult(initial, m)[0]) % MOD
Advertisement
Add Comment
Please, Sign In to add comment
Advertisement