C++ program to count in how many ways we can paint blocks with two conditions

Suppose we have three numbers N, M and K. Consider there are N blocks they are arranged in a row. We consider following two ways to paint them. The paints of two blocks different if and only if the blocks are painted in different colors in following two ways −

• For each block, use one of the M colors to paint it. (it is not mandatory to use all colors)

• There may be at most K pairs of adjacent blocks that are painted in the same color.

If the answer is too large, return result mod 998244353.

So, if the input is like N = 3; M = 2; K = 1, then the output will be 6, because we can paint them in these different formats 112, 121, 122, 211, 212, and 221.

Steps

To solve this, we will follow these steps −

maxm := 2^6 + 5
p := 998244353
Define two large arrays fac and inv or size maxm
Define a function ppow(), this will take a, b, p,
ans := 1 mod p
a := a mod p
while b is non-zero, do:
   if b is odd, then:
      ans := ans * a mod p
   a := a * a mod p
   b := b/2
return ans
Define a function C(), this will take n, m,
if m < 0 or m > n, then:
   return 0
return fac[n] * inv[m] mod p * inv[n - m] mod p
From the main method, do the following
fac[0] := 1
for initialize i := 1, when i < maxm, update (increase i by 1), do:
   fac[i] := fac[i - 1] * i mod p
inv[maxm - 1] := ppow(fac[maxm - 1], p - 2, p)
for initialize i := maxm - 2, when i >= 0, update (decrease i by 1), do:
   inv[i] := (i + 1) * inv[i + 1] mod p
ans := 0
for initialize i := 0, when i <= k, update (increase i by 1), do:
   t := C(n - 1, i)
   tt := m * ppow(m - 1, n - i - 1, p)
   ans := (ans + t * tt mod p) mod p
return ans

Example

Let us see the following implementation to get better understanding −

#include <bits/stdc++.h>
using namespace std;

const long maxm = 2e6 + 5;
const long p = 998244353;
long fac[maxm], inv[maxm];

long ppow(long a, long b, long p){
   long ans = 1 % p;
   a %= p;
   while (b){
      if (b & 1)
         ans = ans * a % p;
      a = a * a % p;
      b >>= 1;
   }
   return ans;
}
long C(long n, long m){
   if (m < 0 || m > n)
      return 0;
   return fac[n] * inv[m] % p * inv[n - m] % p;
}
long solve(long n, long m, long k){
   fac[0] = 1;
   for (long i = 1; i < maxm; i++)
      fac[i] = fac[i - 1] * i % p;
   inv[maxm - 1] = ppow(fac[maxm - 1], p - 2, p);
   for (long i = maxm - 2; i >= 0; i--)
      inv[i] = (i + 1) * inv[i + 1] % p;
   long ans = 0;
   for (long i = 0; i <= k; i++){
      long t = C(n - 1, i);
      long tt = m * ppow(m - 1, n - i - 1, p) % p;
      ans = (ans + t * tt % p) % p;
   }
   return ans;
}
int main(){
   int N = 3;
   int M = 2;
   int K = 1;
   cout << solve(N, M, K) << endl;
}

Input

3, 2, 1

Output

6

Arnab Chakraborty

Updated on 03-Mar-2022 07:00:15

• Related Questions & Answers
• Program to count how many ways we can divide the tree into two trees in Python
• C++ program to count in how many ways we can minimize cable length to connect computers
• Program to count how many ways we can cut the matrix into k pieces in python
• Program to find how many ways we can climb stairs in Python
• In how many ways we can concatenate Strings in Java?
• In how many ways can we split a string in JavaScript?
• Program to count how many blocks are covered k times by walking in Python
• Program to check how many ways we can choose empty cells of a matrix in python
• C++ program to count how many ways two players win or make draw in dice throwing game
• How many ways can we read data from the keyboard in Java?
• C++ program to count at least how many operations needed for given conditions
• Ways to paint stairs with two colors such that two adjacent are not yellow in C++
• Program to count how many times we can find "pizza" with given string characters in Python
• Program to count number of ways we can throw n dices in Python
• What are JSP declarations? In how many ways we can write JSP declarations?