#1083882 (C++17) No.2004 Incremental Coins

• 問題: No.2004 Incremental Coins
• ユーザー: qwewe
• 提出日時: 2025-05-14 13:06:56
• 言語: C++17 (gcc 13.3.0 + boost 1.87.0)
• 結果: TLE
• 実行時間: -
• コード長: 7,101 bytes
• コンパイル時間: 928 ms
• コンパイル使用メモリ: 86,080 KB
• 実行使用メモリ: 30,692 KB
• 最終ジャッジ日時: 2025-05-14 13:08:36
• 合計ジャッジ時間: 5,404 ms
• ジャッジサーバーID （参考情報）: judge3 / judge4

ファイルパターン	結果
sample	AC * 2
other	AC * 6 TLE * 1 -- * 13

ソースコード

#include <iostream>
#include <vector>
#include <numeric>
#include <algorithm>

// Use long long for potentially large numbers like K and coin counts
using ll = long long;

// Problem constraints and inputs
ll N_ll; // N can be up to 2e5
ll K;    // K can be up to 10^18
std::vector<ll> A; // Initial coin counts
std::vector<int> P; // Parent array P_j < j
int N; // Use int for array sizes/indices since N <= 2e5 fits int

const ll MOD = 998244353; // The modulus

// Modular exponentiation: computes (base^exp) % MOD
ll power(ll base, ll exp) {
    ll res = 1;
    base %= MOD;
    while (exp > 0) {
        if (exp % 2 == 1) res = (res * base) % MOD;
        base = (base * base) % MOD;
        exp /= 2;
    }
    return res;
}

// Modular inverse using Fermat's Little Theorem: computes n^(-1) % MOD
// Requires MOD to be prime, which 998244353 is.
ll modInverse(ll n) {
    // Check if n is divisible by MOD. Inverse does not exist.
    // In this problem context, denominators (p! and p) for p <= N < MOD, so this shouldn't happen.
    if (n % MOD == 0) {
        // According to problem constraints this path should not be reached.
        // If it were, we might return 0 or handle error.
        return 0; 
    }
    return power(n, MOD - 2);
}

// Precomputed modular inverses of numbers up to N
std::vector<ll> invNum;

// Precompute modular inverses up to max_val
void precompute_inverses(int max_val) {
    if (max_val < 0) max_val = 0;
    invNum.resize(max_val + 1);
    // Modular inverse for 1 is 1
    if (max_val > 0) invNum[1] = 1; 

    for (int i = 2; i <= max_val; ++i) {
        // Compute modular inverse using the property: inv(i) = -(MOD/i) * inv(MOD % i) % MOD
        // Ensure positive result by adding MOD
        invNum[i] = MOD - (MOD / i) * invNum[MOD % i] % MOD;
    }
}

// Stores binomial coefficients Binom(K, p) mod MOD
std::vector<ll> K_choose_p_vals;

// Compute Binom(K, p) for p = 0..max_p efficiently
// Uses precomputed modular inverses invNum[p]
void compute_K_choose_p(int max_p) {
     if (max_p < 0) return;
    K_choose_p_vals.resize(max_p + 1);
    K_choose_p_vals[0] = 1; // Binom(K, 0) = 1
    
    for (int p = 1; p <= max_p; ++p) {
        // Use formula: Binom(K, p) = Binom(K, p-1) * (K - (p-1)) / p
        
        // Calculate (K - (p-1)) % MOD. Handle potential negativity.
        ll K_minus_p_plus_1 = (K - (p - 1));
        K_minus_p_plus_1 %= MOD;
        if (K_minus_p_plus_1 < 0) K_minus_p_plus_1 += MOD;
        
        ll term1 = K_minus_p_plus_1;
        ll term2 = invNum[p]; // Use precomputed modular inverse of p

        // Calculate Binom(K, p) based on Binom(K, p-1)
        K_choose_p_vals[p] = (K_choose_p_vals[p - 1] * term1) % MOD;
        K_choose_p_vals[p] = (K_choose_p_vals[p] * term2) % MOD;
    }
    // Note: The formula works correctly. If K < p, eventually the term (K - i) for i=K becomes 0,
    // making Binom(K, p) = 0 for all subsequent p.
}

// Adjacency list for the reversed dependency graph:
// children[i] stores list of indices k such that P[k] = i.
// This represents boxes k that send coins to box i.
std::vector<std::vector<int>> children; 

int main() {
    // Use faster I/O operations
    std::ios_base::sync_with_stdio(false);
    std::cin.tie(NULL);

    // Read inputs N and K
    std::cin >> N_ll >> K;
    N = (int)N_ll; // Cast N to int for typical usage like vector sizes
    
    // Resize vectors based on N
    A.resize(N + 1);
    P.resize(N + 1); // P needs indices 1..N
    children.resize(N + 1); // children needs indices 0..N

    // Read initial coin counts A_0 to A_N
    for (int i = 0; i <= N; ++i) {
        std::cin >> A[i];
    }
    // Read parent pointers P_1 to P_N and build the children list
    for (int j = 1; j <= N; ++j) {
        std::cin >> P[j];
         // Problem constraints guarantee P[j] < j. Check P[j] bounds just in case.
         if (P[j] >= 0 && P[j] < j) { 
             children[P[j]].push_back(j);
         }
    }

    // The core logic relies on the fact that the state update A^{(k+1)} = M A^{(k)}
    // is linear, and the matrix M = I + J where J is strictly upper triangular (nilpotent).
    // The final state is A^{(K)} = M^K A^{(0)} = (I+J)^K A^{(0)} = sum_{p=0}^{N} Binom(K, p) J^p A^{(0)}.
    // J^p A^{(0)} becomes zero vector for p > d_max, where d_max is the maximum dependency depth.
    // We can compute this sum iteratively. Let v_p = J^p A^{(0)}. Then v_0 = A^{(0)} and v_p = J v_{p-1}.
    // The final answer is sum_{p=0}^{d_max} Binom(K, p) v_p.

    // Max possible degree is N. Precompute necessary values up to N.
    int max_deg_bound = N; 

    precompute_inverses(max_deg_bound); // Compute modular inverses up to N
    compute_K_choose_p(max_deg_bound); // Compute Binom(K, p) up to N

    // final_A vector stores the result. Initialize with p=0 term: Binom(K,0) * J^0 A^{(0)} = 1 * I * A^{(0)} = A^{(0)}
    std::vector<ll> final_A = A; 
    
    // v_curr stores J^{p-1} * A^{(0)} at the start of iteration p
    std::vector<ll> v_curr = A; // Initially p=1, so this is J^0 A^{(0)}
    std::vector<ll> v_next(N + 1); // Temporary vector to compute J^p A^{(0)}
    
    bool v_curr_is_zero; // Optimization flag to check if J^p A^{(0)} becomes zero

    // Loop for p from 1 up to N (max possible degree)
    for (int p = 1; p <= max_deg_bound; ++p) {
        
        // Compute v_next = J * v_curr
        // The action (J v)_i = sum_{k: P_k=i} v_k, using the children list implementation.
        std::fill(v_next.begin(), v_next.end(), 0);
        v_curr_is_zero = true; // Assume next vector will be zero unless proven otherwise
        for (int i = 0; i <= N; ++i) {
            for (int k : children[i]) { 
                 // k is an index such that P[k] = i. Problem guarantees k > P[k] = i.
                 v_next[i] = (v_next[i] + v_curr[k]) % MOD;
            }
             if (v_next[i] != 0) v_curr_is_zero = false; // Found a non-zero component
        }
        
        // Update v_curr to hold J^p A^{(0)}
        v_curr = v_next; 
        
        // Get precomputed Binom(K, p)
        ll K_p = K_choose_p_vals[p];
        
        // If K_p is 0, this term's contribution is 0. We can skip the addition.
        // Note: This doesn't mean we can break the loop, as maybe future v_p could be nonzero if K_p was zero due to MOD arithmetic specifics.
        // The primary optimization is checking if v_curr (J^p A^{(0)}) is zero.
        if (K_p != 0) {
             // Add the term Binom(K, p) * v_curr to final_A
            for (int i = 0; i <= N; ++i) {
                 final_A[i] = (final_A[i] + K_p * v_curr[i]) % MOD;
                 // Ensure result remains non-negative modulo MOD
                 if (final_A[i] < 0) final_A[i] += MOD; 
            }
        }
        
        // Optimization check: If current vector J^p A^{(0)} is zero, all subsequent J^{p'} A^{(0)} will be zero.
        if (v_curr_is_zero) break; 
    }

    // Output the final coin counts for each box
    for (int i = 0; i <= N; ++i) {
        std::cout << final_A[i] << std::endl;
    }

    return 0;
}