1666F - Fancy Stack - CodeForces Solution


combinatorics dp *2200

Please click on ads to support us..

Python Code:

import os
import sys
from io import BytesIO, IOBase

_str = str
str = lambda x=b"": x if type(x) is bytes else _str(x).encode()
BUFSIZE = 8192

class FastIO(IOBase):
    newlines = 0

    def __init__(self, file):
        self._fd = file.fileno()
        self.buffer = BytesIO()
        self.writable = "x" in file.mode or "r" not in file.mode
        self.write = self.buffer.write if self.writable else None

    def read(self):
        while True:
            b = os.read(self._fd, max(os.fstat(self._fd).st_size, BUFSIZE))
            if not b:
                break
            ptr = self.buffer.tell()
            self.buffer.seek(0, 2), self.buffer.write(b), self.buffer.seek(ptr)
        self.newlines = 0
        return self.buffer.read()

    def readline(self):
        while self.newlines == 0:
            b = os.read(self._fd, max(os.fstat(self._fd).st_size, BUFSIZE))
            self.newlines = b.count(b"\n") + (not b)
            ptr = self.buffer.tell()
            self.buffer.seek(0, 2), self.buffer.write(b), self.buffer.seek(ptr)
        self.newlines -= 1
        return self.buffer.readline()

    def flush(self):
        if self.writable:
            os.write(self._fd, self.buffer.getvalue())
            self.buffer.truncate(0), self.buffer.seek(0)

class IOWrapper(IOBase):
    def __init__(self, file):
        self.buffer = FastIO(file)
        self.flush = self.buffer.flush
        self.writable = self.buffer.writable
        self.write = lambda s: self.buffer.write(s.encode("ascii"))
        self.read = lambda: self.buffer.read().decode("ascii")
        self.readline = lambda: self.buffer.readline().decode("ascii")

sys.stdin, sys.stdout = IOWrapper(sys.stdin), IOWrapper(sys.stdout)
input = lambda: sys.stdin.readline().rstrip("\r\n")

from collections import Counter

def main():
    t = int(input())
    while t:
        n = int(input())
        a = list(map(int, input().strip().split(' ')))
        MOD = 998244353
        fac = [1] * (n + 1)
        inv = [1] * (n + 1)
        finv = [1] * (n + 1)
        for i in range(2, n + 1):
            fac[i] = fac[i-1] * i % MOD
            inv[i] = MOD - MOD // i * inv[MOD%i] % MOD
            finv[i] = finv[i-1] * inv[i] % MOD
        
        def C(n, m):
            if m == 0:
                return 1
            if m > n:
                return 0
            return fac[n] * finv[m] * finv[n-m] % MOD

        cnt = Counter(a)
        cnt = sorted(cnt.items(), key=lambda x: x[0], reverse=True)
        dp = [[0] * (n//2+1) for _ in range(len(cnt)+1)]
        dp[0][0] = 1
        pre_sum = [0] * (len(cnt)+1)
        for i in range(1, len(cnt)+1):
            pre_sum[i] = pre_sum[i-1] + cnt[i-1][1]
        for i in range(len(cnt)):
            for j in range(min(i+1, n//2+1)):
                small = pre_sum[i]-j
                if j < n // 2:
                    dp[i+1][j+1] = (dp[i+1][j+1] + dp[i][j] * C(j-1-small, cnt[i][1]-1)) % MOD
                slot = j-1-small
                if j == n // 2:
                    slot += 1
                dp[i+1][j] = (dp[i+1][j] + dp[i][j] * C(slot, cnt[i][1])) % MOD
        print(dp[len(cnt)][n//2])
        t -= 1
    return

main()

C++ Code:

/******************************************************************************

                              Online C++ Compiler.
               Code, Compile, Run and Debug C++ program online.
Write your code in this editor and press "Run" button to compile and execute it.

*******************************************************************************/

#include <iostream>

using namespace std;
using ll = long long;

int T;
const int MOD = 998244353;
ll a[5001];
ll dp[5001][2501], fac[2501], inv[2501];

ll modpow(ll x, ll n){
    ll a = 1;
    x %= MOD;
    while(n>0){
        if(n%2==1){
            a = (a*x) % MOD;
        }
        x = (x*x) % MOD;
        n/=2;
    }
    return a;
}

void factorial(){
    fac[0] = 1;
    for(int i=1;i<=2500;i++){
        fac[i] = (fac[i-1] * i)%MOD;
    }
}

void inverse(){
    inv[2500] = modpow(fac[2500],MOD-2);
    for(int i=2500;i>0;i--){
        inv[i-1] = (inv[i] * i)%MOD;
    }
}

ll nCr(ll n, ll r){
    if(r == 0) return 1;
    if(n < r) return 0;
    return (fac[n] * inv[r] % MOD) * inv[n-r]%MOD;
}

int main(){
    ios::sync_with_stdio(false);
    cin.tie(nullptr);
    
    factorial(); inverse();
    
    cin >> T;
    while(T--){
        int N; cin >> N;
        for(int i=N-1;i>=0;i--){
            cin >> a[i];
        }
        
        dp[0][0] = 1;
        for(int i=0;i<N;i++){
            ll cnt = 1, t = i;
            while(t!=N-1 && a[t] == a[t + 1]){
                cnt++; t++;
            }

            for(int j=0;j<=N/2;j++){
                if(j > i) break;

                if(j == N/2){
                    dp[i + cnt][j] += dp[i][j] * nCr(N/2 - (i-j), cnt) % MOD;
                } else {
                    dp[i + cnt][j+1] += dp[i][j] * nCr(j-1 - (i-j),cnt-1) % MOD;
                    dp[i + cnt][j] += dp[i][j] * nCr(j-1 - (i-j),cnt) % MOD;
                    
                    dp[i + cnt][j+1] %= MOD;
                }
                dp[i + cnt][j] %= MOD;
            }
            
            i = t;
        }
        
        cout << dp[N][N/2] << "\n";
        
        for(int i=0;i<=N;i++){
            for(int j=0;j<=N/2;j++){
                dp[i][j] = 0;
            }
        }
    }
    return 0;
}


Comments

Submit
0 Comments
More Questions

1665D - GCD Guess
29A - Spit Problem
1097B - Petr and a Combination Lock
92A - Chips
1665B - Array Cloning Technique
1665A - GCD vs LCM
118D - Caesar's Legions
1598A - Computer Game
1605A - AM Deviation
1461A - String Generation
1585B - Array Eversion
1661C - Water the Trees
1459A - Red-Blue Shuffle
1661B - Getting Zero
1661A - Array Balancing
1649B - Game of Ball Passing
572A - Arrays
1455A - Strange Functions
1566B - MIN-MEX Cut
678C - Joty and Chocolate
1352E - Special Elements
1520E - Arranging The Sheep
1157E - Minimum Array
1661D - Progressions Covering
262A - Roma and Lucky Numbers
1634B - Fortune Telling
1358A - Park Lighting
253C - Text Editor
365B - The Fibonacci Segment
75A - Life Without Zeros