C++ Program to find pairs of sequences where sequence holds minimum and maximum elements

Suppose we have three numbers N, M and K. There are N horizontal rows and M vertical rows. We shall write an integer in between 1 and K on each cell, and define sequences A and B, such that −

for each i in range 1 to N, A[i] is minimum of all elements in ith row
for each j in range 1 to M, B[j] is maximum of all elements in jth column

We have to find the number of pairs (A, B). If the answer is too large, return result mod 998244353.

So, if the input is like N = 2; M = 2; K = 2, then the output will be 7, because the (A[1], A[2], B[1], B[2]) are (1,1,1,1), (1,1,1,2), (1,1,2,1), (1,1,2,2), (1,2,2,2), (2,1,2,2), or (2,2,2,2).

Steps

To solve this, we will follow these steps −

p := 998244353
Define a function power(), this will take a, b, and return (a^b) mod p
From the main method, do the following:
if n is same as 1, then:
   return power(K, m)
if m is same as 1, then:
   return power(K, n)
ans := 0
for initialize t := 1, when t <= K, update (increase t by 1), do:
   ans := (ans + (power(t, n) - power(t - 1, n) + p) mod p * power(K - t + 1, m)) mod p
return ans

Example

Let us see the following implementation to get better understanding −

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

long p = 998244353;

long power(long a, long b, long ret = 1){
   for (; b; b >>= 1, a = a * a % p)
      if (b & 1)
         ret = ret * a % p;
   return ret;
}
long solve(int n, int m, int K){
   if (n == 1)
      return power(K, m);
   if (m == 1)
      return power(K, n);
   long ans = 0;
   for (long t = 1; t <= K; t++){
      ans = (ans + (power(t, n) - power(t - 1, n) + p) % p * power(K - t + 1, m)) % p;
   }
   return ans;
}
int main(){
   int N = 2;
   int M = 2;
   int K = 2;
   cout << solve(N, M, K) << endl;
}