Codeforces 1172C2 Nauuo and Pictures (hard version) dp

Nauuo and Pictures (hard version

First, consider the simple version, seek out one by one, dp, divided into three lump, lump currently required photo, cook in addition to the current favorite photos, cook in addition to the current nasty pictures.

Single dp 50 ^ 4

Feel hard is also quite simple. . 

Finally, we first calculate the total favorite photos w, w and hate the total picture, and then each contribution is in proportion in the original w.

#include<bits/stdc++.h>
#define LL long long
#define LD long double
#define ull unsigned long long
#define fi first
#define se second
#define mk make_pair
#define PLL pair<LL, LL>
#define PLI pair<LL, int>
#define PII pair<int, int>
#define SZ(x) ((int)x.size())
#define ALL(x) (x).begin(), (x).end()
#define fio ios::sync_with_stdio(false); cin.tie(0);

using namespace std;

const int N = 3000 + 7;
const int inf = 0x3f3f3f3f;
const LL INF = 0x3f3f3f3f3f3f3f3f;
const int mod = 998244353;
const double eps = 1e-10;
const double PI = acos(-1);

template<class T, class S> inline void add(T &a, S b) {a += b; if(a >= mod) a -= mod;}
template<class T, class S> inline void sub(T &a, S b) {a -= b; if(a < 0) a += mod;}
template<class T, class S> inline bool chkmax(T &a, S b) {return a < b ? a = b, true : false;}
template<class T, class S> inline bool chkmin(T &a, S b) {return a > b ? a = b, true : false;}

mt19937 rng(chrono::steady_clock::now().time_since_epoch().count());

int power(int a, int b) {
    int ans = 1;
    while(b) {
        if(b & 1) ans = 1LL * ans * a % mod;
        a = 1LL * a * a % mod; b >>= 1;
    }
    return ans;
}

int n, m, a[200007], w[200007], r[200007];
int sumL, sumH, EL, EH;

int dp[N][N];

int main() {

    scanf("%d%d", &n, &m);
    for(int i = 1; i <= n; i++) scanf("%d", &a[i]);
    for(int i = 1; i <= n; i++) {
        scanf("%d", &w[i]);
        if(a[i]) sumL += w[i];
        else sumH += w[i];
    }

    EL = sumL;
    EH = sumH;

    dp[0][0] = 1;

    for(int i = 0; i < m; i++) {
        for(int a = 0; a <= i; a++) {
            int b = i - a;
            add(dp[i + 1][a + 1], 1LL * dp[i][a] * (sumL + a) % mod * power(sumL + sumH + a - b, mod - 2) % mod);
            add(dp[i + 1][a], 1LL * dp[i][a] * (sumH - b) % mod * power(sumL + sumH + a - b, mod - 2) % mod);
        }
    }

    for(int a = 0; a <= m; a++) {
        int b = m - a;
        add(EL, 1LL * a * dp[m][a] % mod);
        sub(EH, 1LL * b * dp[m][a] % mod);
    }

    for(int i = 1; i <= n; i++) {
        if(a[i]) {

            r[i] = 1LL * EL * w[i] % mod * power(sumL, mod - 2) % mod;

        } else {

            r[i] = 1LL * EH * w[i] % mod * power(sumH, mod - 2) % mod;

        }
    }

    for(int i = 1; i <= n; i++) printf("%d\n", r[i]);
    return 0;
}

/*
*/