1725C - Circular Mirror - CodeForces Solution

#include<bits/stdc++.h>
using namespace std;
#define ll long long 
#define pb push_back
#define ub upper_bound
#define lb lower_bound
#define mp make_pair
#define pii pair<ll int,ll int>
#define umap unordered_map
#define popcount(x) __builtin_popcountll(x)
#define all(v) v.begin() , v.end()
#define PI 3.141592653589793238
#define E 2.7182818284590452353602874713527
#define M 998244353
const long long INF = 1e18;
#define error(args...) { string _s = #args; replace(_s.begin(), _s.end(), ',', ' '); stringstream _ss(_s); istream_iterator<string> _it(_ss); err(_it, args); }

void err(istream_iterator<string> it) {}
template<typename T, typename... Args>
void err(istream_iterator<string> it, T a, Args... args) {
  cerr << *it << " = " << a << endl;
  err(++it, args...);
}

ll int fact[300001], invfact[300001];

ll power(ll a, ll b, const ll mod){ 
    ll x = 1, y = a; 
    while (b){ 
        if (b & 1){ 
            x *= y;
            x %= mod; 
        } 
        y *= y;
        y %= mod; 
        b >>= 1; 
    } 
    return x; 
}

ll modular_inverse(ll n, const ll mod){ 
    return power(n, mod-2, mod); 
} 
 
ll nCr(ll n, ll k, ll mod){ 
    return (fact[n]*((invfact[k]*invfact[n-k])%mod))%mod;
} 


int main()
{
ios_base::sync_with_stdio(0);cin.tie(0);cout.tie(0);
// freopen("input.txt", "r", stdin);
// freopen("output.txt", "w", stdout);
fact[0] = 1;
for (int i = 1; i <= 300000; i++) {
        fact[i] = fact[i - 1] * i;
        fact[i] %= M;
}

invfact[300000] = modular_inverse(fact[300000], M);
for (int i = 299999; i >= 0; i--) {
        invfact[i] = invfact[i + 1] * (i + 1) % M;
}
ll int n, m;
cin >> n >> m;
ll int d[n + 1] = {};
map<ll int,ll int>p;
for (int i = 1; i <=n; i++)
{
  cin >> d[i];
  d[i] += d[i - 1];
  p[d[i]]=1;
}
if (d[n] % 2 == 1) {
  cout << power(m, n, M) << '\n';
  return 0;
}
ll int x = d[n] / 2,z=0;
for (int i = 1; i<=n; i++)
{
  if (p[d[i] + x] == 1)
    z++;
}
// error(n,z);
ll int ans=0;
for (int i = 0; i <= z; i++) {
  if (m < i)
    break;
  ll int y = nCr(z, i, M);
  y *= nCr(m, i, M);
  y %= M;
  y*=fact[i];
  y %= M;
  ll int lf = z - i, lc = m - i;
  // error(lf,lc,y);
  y *= power((lc * (lc - 1))%M, lf, M);
  y %= M;
  // error(y);
  y *= power(lc, n - (2 * z), M);
  y %= M;
  // y*=4;  
  // error(power(lc, n - (2 * z), M));
  // error(y);
  ans += y;
  ans%=M;
}
cout<<ans;
}