diff --git a/src/problems/3405-count-the-number-of-arrays-with-k-matching-adjacent-elements.cpp b/src/problems/3405-count-the-number-of-arrays-with-k-matching-adjacent-elements.cpp new file mode 100644 index 0000000..82d438d --- /dev/null +++ b/src/problems/3405-count-the-number-of-arrays-with-k-matching-adjacent-elements.cpp @@ -0,0 +1,80 @@ +/** +* [3405] Count the Number of Arrays with K Matching Adjacent Elements + */ +#include +#include +using namespace std; + + +// submission codes start here + +constexpr long long MOD = 1e9 + 7; +constexpr long long UPPER_BOUND = 1e5; + +static long long fact[UPPER_BOUND]; +static long long inverseFact[UPPER_BOUND]; + +class Solution +{ + static long long quickPower(long long x, int n) + { + long long result = 1; + + while (n > 0) + { + if ((n & 1) == 1) + { + result = result * x % MOD; + } + + x = x * x % MOD; + n >>= 1; + } + + return result; + } + + static long long combine(int n, int m) + { + return fact[n] * inverseFact[m] % MOD * inverseFact[n - m] % MOD; + } + + static void init() + { + if (fact[0] != 0) + { + return; + } + + fact[0] = 1; + for (int i = 1; i < UPPER_BOUND; ++i) + { + fact[i] = fact[i - 1] * i % MOD; + } + + // Modular Multiplicative Inverse is calculated by the quick power. + inverseFact[UPPER_BOUND - 1] = quickPower(fact[UPPER_BOUND - 1], MOD - 2); + for (int i = UPPER_BOUND - 1; i > 0; --i) + { + inverseFact[i - 1] = inverseFact[i] * i % MOD; + } + } + +public: + int countGoodArrays(int n, int m, int k) + { + init(); + const long long result = combine(n - 1, k) * m % MOD * quickPower(m - 1, n - k - 1) % MOD; + return result; + } +}; + +// submission codes end + +TEST(P3405, Test1) +{ + Solution s; + ASSERT_EQ(4, s.countGoodArrays(3, 2, 1)); + ASSERT_EQ(6, s.countGoodArrays(4, 2, 2)); + ASSERT_EQ(2, s.countGoodArrays(5, 2, 0)); +} \ No newline at end of file