Data Structure
Networking
RDBMS
Operating System
Java
MS Excel
iOS
HTML
CSS
Android
Python
C Programming
C++
C#
MongoDB
MySQL
Javascript
PHP
Physics
Chemistry
Biology
Mathematics
English
Economics
Psychology
Social Studies
Fashion Studies
Legal Studies
- Selected Reading
- UPSC IAS Exams Notes
- Developer's Best Practices
- Questions and Answers
- Effective Resume Writing
- HR Interview Questions
- Computer Glossary
- Who is Who
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
234 Views
