A bunch new problems

Codeforces Round 856 (Div. 2)/D. Counting Factorizations.cpp

@@ -78,20 +78,17 @@ auto operator>>(istream  &is, vector<vector<v>> &vec)->istream& {
 const int oo = INT32_MAX >> 1;
 const ll OO = INT64_MAX >> 1;
 
-const ll mod = 998244353;
-
-constexpr MAXN = 1e6 + 10;
+constexpr int MAXN = 1e6 + 10;
 bool prime[MAXN];
-bool vis[MAXN];
-
-vl fact[MAXN];
-vl inv[MAXN];
+ll fact[MAXN];
+ll inv[MAXN];
+ll finv[MAXN];
 
+ll mod = 998244353;
 
 void pre()
 {
     fill(prime, prime + MAXN, true);
-    prime[0] = false;
     prime[1] = false;
 
     nrep(i, 2, MAXN) {
@@ -105,50 +102,38 @@ void pre()
     }
 
     fact[0] = 1;
-    fact[1] = 1;
-    nrep(i, 2, MAXN) {
+    nrep(i, 1, MAXN) {
         fact[i] = fact[i - 1] * i % mod;
     }
 
-    inv[0] = 1;
     inv[1] = 1;
     nrep(i, 2, MAXN) {
         inv[i] = (mod - mod / i) * inv[mod % i] % mod;
     }
 
-    nrep(i, 2, MAXN) {
-        inv[i] = inv[i - 1] * inv[i] % mod;
+    finv[0] = 1;
+    nrep(i, 1, MAXN) {
+        finv[i] = finv[i - 1] * inv[i] % mod;
     }
 }
 
-#define TEST 1
+#define TEST 0
 
 void solve()
 {
     int n;
     cin >> n;
+    vi a(2 * n);
+    cin >> a;
 
-    vi act(n << 1);
-    cin >> act;
-
-    set<int> up;
-
-    vi ot;
     vi primes;
-
-    rep(i, n << 1) {
-        if (prime[act[i]]) {
-            if (vis[act[i]]) {
-                ot.push_back(act[i]);
-            } else {
-                up.insert(act[i]);
-                vis[act[i]] = true;
-                primes.push_back(act[i]);
-            }
-            continue;
+    map<int, ll> c;
+    rep(i, 2 * n) {
+        if (prime[a[i]] && !c.count(a[i])) {
+            primes.push_back(a[i]);
         }
 
-        ot.push_back(act[i]);
+        c[a[i]]++;
     }
 
     if (primes.size() < n) {
@@ -156,38 +141,27 @@ void solve()
         return;
     }
 
-    auto comb = [&](int n, int k) {
-        return fact[n] * inv[k] % mod * inv[n - k] % mod;
-    };
-
-    ll ans = comb(primes.size(), n);
-
-    if (ot.empty()) {
-        cout << ans << '\n';
-        return;
-    }
-
-    sortv(ot);
-
-    int c = 1;
-    vi tmp;
-
-    nrep(i, 1, ot.size()) {
-        if (ot[i] == ot[i - 1]) {
-            c++;
-            continue;
+    ll non = 1;
+    repv(i, c) {
+        if (!prime[i.first]) {
+            non = (non * finv[i.second]) % mod;
         }
-
-        tmp.push_back(c);
-        c = 1;
     }
-    tmp.push_back(c);
 
+    vvl dp(primes.size() + 1, vl(n + 1));
+    dp[0][0] = fact[n] * non % mod;
 
+    nrep(i, 1, primes.size() + 1) {
+        int cur = primes[i - 1];
+        ll co = c[cur];
+        dp[i][0] = dp[i - 1][0] * finv[co] % mod;
 
-    repv(i, up) {
-        vis[i] = false;
+        nrep(j, 1, n + 1) {
+            dp[i][j] = (dp[i - 1][j] * finv[co] % mod + dp[i - 1][j - 1] * finv[co - 1] % mod) % mod;
+        }
     }
+
+    cout << dp[primes.size()][n] << '\n';
 }
 
 int main()