Create count-the-number-of-arrays-with-k-matching-adjacent-elements.py

Author: kamyu104

Python/count-the-number-of-arrays-with-k-matching-adjacent-elements.py

@@ -0,0 +1,23 @@
+# Time:  O(n + logm)
+# Space: O(n)
+
+# combinatorics, fast exponentiation
+MOD = 10**9+7
+FACT, INV, INV_FACT = [[1]*2 for _ in xrange(3)]
+def nCr(n, k):
+    while len(INV) <= n:  # lazy initialization
+        FACT.append(FACT[-1]*len(INV) % MOD)
+        INV.append(INV[MOD%len(INV)]*(MOD-MOD//len(INV)) % MOD)  # https://cp-algorithms.com/algebra/module-inverse.html
+        INV_FACT.append(INV_FACT[-1]*INV[-1] % MOD)
+    return (FACT[n]*INV_FACT[n-k] % MOD) * INV_FACT[k] % MOD
+
+
+class Solution(object):
+    def countGoodArrays(self, n, m, k):
+        """
+        :type n: int
+        :type m: int
+        :type k: int
+        :rtype: int
+        """
+        return (nCr(n-1, k)*(m*pow(m-1, (n-1)-k, MOD)))%MOD