#1140903 (Rust) No.3414 Aperiodic Sequence

• 問題: No.3414 Aperiodic Sequence
• ユーザー: wolgnik
• 提出日時: 2025-12-18 13:59:57
• 言語: Rust (1.83.0 + proconio)
• 結果: WA
• 実行時間: -
• コード長: 4,986 bytes
• コンパイル時間: 14,822 ms
• コンパイル使用メモリ: 403,572 KB
• 実行使用メモリ: 16,076 KB
• 最終ジャッジ日時: 2025-12-20 23:31:16
• 合計ジャッジ時間: 19,942 ms
• ジャッジサーバーID （参考情報）: judge5 / judge2

ファイルパターン	結果
sample	WA * 2
other	WA * 5 TLE * 1 -- * 9

ソースコード

/*
 * 愚直に k 乗和をやる
 */

use proconio::input;
use std::thread;
use std::sync::Arc;

const MOD: u64 = 998244353;

#[inline(always)]
fn pow_mod(mut base: u64, mut exp: u64) -> u64 {
    let mut result = 1u64;
    while exp > 0 {
        if exp & 1 == 1 {
            result = result * base % MOD;
        }
        exp >>= 1;
        base = base * base % MOD;
    }
    result
}

fn main() {
    input! {
        n: usize,
        m: usize,
    }
    
    let v = vec![100; n];

    let v_mod: Vec<u64> = v.iter().map(|&x| x % MOD).collect();

    // スレッド数
    let num_threads = 8usize;

    let g: Vec<u64> = if n == 0 {
        vec![0; m + 1]
    } else {
        let v_mod = Arc::new(v_mod);
        let chunk_size = (n + num_threads - 1) / num_threads;
        let m_copy = m;  // 明示的にコピー

        let handles: Vec<_> = (0..num_threads)
            .filter_map(|t| {
                let start = t * chunk_size;
                let end = ((t + 1) * chunk_size).min(n);
                if start >= end {
                    return None;
                }
                
                let v_mod = Arc::clone(&v_mod);

                Some(thread::spawn(move || {
                    let mut local_g = vec![0u64; m_copy + 1];
                    local_g[0] = (end - start) as u64;

                    for i in start..end {
                        let vi = v_mod[i];
                        let mut p = vi;
                        for k in 1..=m_copy {
                            local_g[k] = (local_g[k] + p) % MOD;
                            p = p * vi % MOD;
                        }
                    }
                    local_g
                }))
            })
            .collect();

        // マージ
        let mut g = vec![0u64; m + 1];
        for handle in handles {
            let local_g = handle.join().unwrap();
            for k in 0..=m {
                g[k] = (g[k] + local_g[k]) % MOD;
            }
        }
        g
    };

    // メビウス関数
    let mut mobius = vec![0i8; m + 1];
    if m >= 1 {
        mobius[1] = 1;
    }
    let mut spf: Vec<u32> = (0..=m as u32).collect();
    let sqrt_m = ((m as f64).sqrt() as usize) + 1;
    for i in 2..=sqrt_m.min(m) {
        if spf[i] == i as u32 {
            let mut j = i * i;
            while j <= m {
                if spf[j] == j as u32 {
                    spf[j] = i as u32;
                }
                j += i;
            }
        }
    }

    for i in 2..=m {
        let p = spf[i] as usize;
        let prev = i / p;
        mobius[i] = if prev % p == 0 { 0 } else { -mobius[prev] };
    }

    // 結果計算（並列）
    let g = Arc::new(g);
    let mobius = Arc::new(mobius);

    let handles: Vec<_> = (0..num_threads)
        .filter_map(|t| {
            let chunk_size = (m + num_threads - 1) / num_threads;
            let start = t * chunk_size + 1;
            let end = ((t + 1) * chunk_size + 1).min(m + 1);
            
            if start >= end {
                return None;
            }

            let g = Arc::clone(&g);
            let mobius = Arc::clone(&mobius);

            Some(thread::spawn(move || {
                let mut local_res = Vec::with_capacity(end - start);
                for k in start..end {
                    let mut f = 0u64;
                    let mut d = 1usize;
                    while d * d <= k {
                        if k % d == 0 {
                            let mu = mobius[k / d];
                            if mu == 1 {
                                f = (f + pow_mod(g[k / d], d as u64)) % MOD;
                            } else if mu == -1 {
                                f = (f + MOD - pow_mod(g[k / d], d as u64) % MOD) % MOD;
                            }
                            if d != k / d {
                                let d2 = k / d;
                                let mu2 = mobius[d];
                                if mu2 == 1 {
                                    f = (f + pow_mod(g[d], d2 as u64)) % MOD;
                                } else if mu2 == -1 {
                                    f = (f + MOD - pow_mod(g[d], d2 as u64) % MOD) % MOD;
                                }
                            }
                        }
                        d += 1;
                    }
                    local_res.push(f);
                }
                (start, local_res)
            }))
        })
        .collect();

    // 結果をマージ
    let mut results: Vec<(usize, Vec<u64>)> = handles
        .into_iter()
        .map(|h| h.join().unwrap())
        .collect();
    results.sort_by_key(|x| x.0);

    let mut out = String::new();
    for (i, (_, chunk)) in results.iter().enumerate() {
        for (j, &val) in chunk.iter().enumerate() {
            if i > 0 || j > 0 {
                out.push(' ');
            }
            out.push_str(&val.to_string());
        }
    }
    println!("{}", out);
}