提出 #21453866

---

ソースコード 拡げる

// https://atcoder.jp/contests/agc052/tasks/agc052_c
//
#![allow(unused_imports)]
use std::io::*;
use std::io::Write;
use std::fmt::*;
use std::str::*;
use std::cmp::*;
use std::collections::*;

// Input macros.
// Original by tanakh: https://qiita.com/tanakh/items/0ba42c7ca36cd29d0ac8
#[allow(unused_macros)]
macro_rules! input {
    (source = $s:expr, $($r:tt)*) => {
        let mut iter = $s.split_whitespace();
        input_inner!{iter, $($r)*}
    };
    ($($r:tt)*) => {
        let s = {
            use std::io::Read;
            let mut s = String::new();
            std::io::stdin().read_to_string(&mut s).unwrap();
            s
        };
        let mut iter = s.split_whitespace();
        input_inner!{iter, $($r)*}
    };
}

#[allow(unused_macros)]
macro_rules! input_inner {
    ($iter:expr) => {};
    ($iter:expr, ) => {};

    ($iter:expr, $var:ident : $t:tt $($r:tt)*) => {
        let $var = read_value!($iter, $t);
        input_inner!{$iter $($r)*}
    };

    ($iter:expr, mut $var:ident : $t:tt $($r:tt)*) => {
        let mut $var = read_value!($iter, $t);
        input_inner!{$iter $($r)*}
    };
}

#[allow(unused_macros)]
macro_rules! read_value {
    ($iter:expr, ( $($t:tt),* )) => {
        ( $(read_value!($iter, $t)),* )
    };

    ($iter:expr, [ $t:tt ; $len:expr ]) => {
        (0..$len).map(|_| read_value!($iter, $t)).collect::<Vec<_>>()
    };

    ($iter:expr, [ next / $t:tt ]) => {
        {
            let len = read_value!($iter, usize);
            (0..len).map(|_| read_value!($iter, $t)).collect::<Vec<_>>()
        }
    };

    ($iter:expr, switch) => {
        {
            let ty = read_value!($iter, i32);
            if ty == 1 {
                vec![ty, read_value!($iter, i32), read_value!($iter, i32)]
            } else if ty == 2 {
                vec![ty, read_value!($iter, i32)]
            } else {
                vec![ty, read_value!($iter, i32)]
            }
        }
    };

    ($iter:expr, chars) => {
        read_value!($iter, String).chars().collect::<Vec<char>>()
    };

    ($iter:expr, usize1) => {
        read_value!($iter, usize) - 1
    };

    ($iter:expr, $t:ty) => {
        $iter.next().unwrap().parse::<$t>().expect("Parse error")
    };
}

#[allow(unused_macros)]
macro_rules! read_line {
    ($t:tt) => {
        {
            let mut s = String::new();
            std::io::stdin().read_line(&mut s).unwrap();
            s.trim_right().parse::<$t>().unwrap()
        }
    }
}

#[allow(unused_macros)]
macro_rules! dvec {
    ($t:expr ; $len:expr) => {
        vec![$t; $len]
    };

    ($t:expr ; $len:expr, $($rest:expr),*) => {
        vec![dvec!($t; $($rest),*); $len]
    };
}

#[allow(unused_macros)]
macro_rules! ifv {
    ($t:expr, $a:expr, $b: expr) => {
        if $t { $a } else { $b }
    }
}

#[allow(unused_macros)]
macro_rules! fill {
    ($t:expr, $v:expr) => {
        for i in 0..$t.len() {
            $t[i] = $v;
        }
    };
}

#[allow(unused_macros)]
macro_rules! join {
    ($t:expr, $glue:expr) => {
        $t.into_iter().map(|w| w.to_string()).collect::<Vec<_>>().join($glue)
    };
}

#[allow(unused_macros)]
macro_rules! debug {
    ($($a:expr),*) => {
        eprintln!(concat!($(stringify!($a), " = {:?}, "),*), $($a),*);
    }
}

// ===


fn next_permutation(a: &mut Vec<usize>) -> bool {
    let n = a.len();
    if n == 1 {
        return false;
    }
    let mut x = n-2;
    while x >= 0 && a[x] >= a[x+1] {
        if x == 0 {
            return false;
        }
        x -= 1;
    }
    let mut y = n-1;
    while y > x && a[y] <= a[x] {
        y -= 1;
    }
    let mut tmp = a[x];
    a[x] = a[y];
    a[y] = tmp;

    // sort
    let mut hp = BinaryHeap::new();
    for i in x+1..n {
        hp.push(a[i]);
    }
    for i in (x+1..n).rev() {
        a[i] = hp.pop().unwrap();
    }
    true
}
//===


fn find(a: &mut Vec<i64>, p: i64) -> bool {
    let n = a.len();
    let mut wo = vec![0; n];
    for i in 0..n {
        wo[i] = i;
    }
    loop {
        let mut psum = 0;
        let mut ok = true;
        for i in 0..n {
            psum += a[wo[i]];
            if psum % p == 0 {
                ok = false;
            }
        }
        if ok {
            return true;
        }
        if !next_permutation(&mut wo) {
            break;
        }
    }
    false
}

fn solve(idx: usize, last: i64, n: usize, p: i64, a: &mut Vec<i64>) -> i64 {
    if idx == n {
        let r = find(a, p);
        let mut wo = 0;
        for i in 0..n {
            wo += a[i];
        }
        if wo % p != 0 && !r {
            debug!(a, r);
        }
        let mut c = dvec![0; n+1, n+1];
        let mut aa = 1;
        for i in 0..=n {
            c[i][0] = 1;
            c[i][i] = 1;
            for j in 1..i {
                c[i][j] = (c[i-1][j-1] + c[i-1][j]) % MOD;
            }
        }
        let mut dig = vec![0; p as usize];
        for i in 0..n {
            dig[a[i] as usize] += 1;
        }
        let mut left = n;
        for i in 0..p {
            aa *= c[left][dig[i as usize]];
            aa %= MOD;
            left -= dig[i as usize];
        }
        return ifv!(r, aa, 0);
    }
    let mut ret = 0;
    for i in last..p {
        a[idx] = i;
        ret += solve(idx+1, i, n, p, a);
        a[idx] = 0;
    }
    ret %= MOD;
    ret
}

fn experiment(n: usize, p: i64) -> i64 {
    let mut a = vec![0; n];
    solve(0, 1, n, p, &mut a)
}

fn ex3(n1: usize, n2: usize, n3: usize, n4: usize, n: usize, p: usize) -> bool {
    let mut dp = dvec![0; n1+1, n2+1, n3+1, n4+1, p];
    dp[n1][n2][n3][n4][0] = 1;
    for i in (0..=n1).rev() {
        for j in (0..=n2).rev() {
            for k in (0..=n3).rev() {
                for l in (0..=n4).rev() {
                    let f = ifv!(i+j+k+l == n, 0, 1);
                    for pi in f..p {
                        let base = dp[i][j][k][l][pi];
                        if base == 0 {
                            continue;
                        }
                        if i >= 1 {
                            dp[i-1][j][k][l][(pi+1)%p] = 1;
                        }
                        if j >= 1 {
                            dp[i][j-1][k][l][(pi+2)%p] = 1;
                        }
                        if k >= 1 {
                            dp[i][j][k-1][l][(pi+3)%p] = 1;
                        }
                        if l >= 1 {
                            dp[i][j][k][l-1][(pi+4)%p] = 1;
                        }
                    }
                }
            }
        }
    }
    let mut no = 0;
    for i in 1..p {
        no = max(no, dp[0][0][0][0][i]);
    }
    return no == 1;
}

fn ex2(n: usize, p: usize) {
    for n1 in 0..n {
        for n2 in 0..n {
            for n3 in 0..n {
                for n4 in 0..n {
                    if n1+n2+n3+n4 != n {
                        continue;
                    }
                    if (n1+n2*2+n3*3+n4*4) % p == 0 {
                        // continue;
                    }
                    if !ex3(n1, n2, n3, n4, n, p) {
                        debug!(n1, n2, n3, n4);
                    }
                }
            }
        }
    }


}

const MOD: i64 = 998244353;

fn main() {
    input! {
        n: usize, p: i64
    };
    if p == 2 {
        println!("{}", ifv!(n == 1, "1", "0"));
        return;
    }
    // debug!(experiment(n, p));
    // ex2(20, 5);
    let mut zero = 1;
    let mut one = 0;
    for i in 0..n {
        let tz = one * (p-1) % MOD;
        let to = (zero + one * (p-2)) % MOD;
        zero = tz;
        one = to;
    }
    let base = one * (p-1) % MOD;
    if n as i64 + 1 <= p {
        println!("{}", base);
        return;
    }
    let mut c = dvec![0; n+1, n+1];
    let mut aa = 1;
    for i in 0..=n {
        c[i][0] = 1;
        c[i][i] = 1;
        for j in 1..i {
            c[i][j] = (c[i-1][j-1] + c[i-1][j]) % MOD;
        }
    }

    let p = p as usize;
    let mut leak = 1;
    let mut dp = dvec![0; n+1, n+1];
    // for i in n/2+1..n {
    //     dp[i][p-1] = c[n][i];
    // }


    let mut psum = vec![0; n+2];
    for i in 1..=n {
        psum[0] = 0;
        for s in 0..=n {
            psum[s+1] = psum[s] + dp[i-1][s];
            psum[s+1] %= MOD;
        }
        for s in p-1..n {
            let f = s-p+1;
            let t = s-2;
            dp[i][s] = (psum[t+1] + MOD - psum[f]) % MOD;
        }
        if i >= n/2+1 && i < n {
            dp[i][p-1] += c[n][i];
            dp[i][p-1] %= MOD;
        }
    }

    for s in 0..n {
        if (n-s)%p == 1 {
            continue;
        }
        leak += dp[n][s];
        leak %= MOD;
    }


    leak *= (p-1) as i64;
    leak %= MOD;

    let mut all = 1;
    for i in 0..n {
        all *= (p-1) as i64;
        all %= MOD;
    }
    // debug!(all, zero, leak);
    println!("{}", (all + MOD + MOD - zero - leak) % MOD);
}

提出情報

コンパイルエラー

warning: unused variable: `i`
   --> src/main.rs:322:9
    |
322 |     for i in 0..n {
    |         ^ help: consider prefixing with an underscore: `_i`
    |
    = note: `#[warn(unused_variables)]` on by default

warning: unused variable: `aa`
   --> src/main.rs:334:13
    |
334 |     let mut aa = 1;
    |             ^^ help: consider prefixing with an underscore: `_aa`

warning: unused variable: `i`
   --> src/main.rs:382:9
    |
382 |     for i in 0..n {
    |         ^ help: consider prefixing with an underscore: `_i`

warning: variable does not need to be mutable
   --> src/main.rs:161:9
    |
161 |     let mut tmp = a[x];
    |         ----^^^
    |         |
    |         help: remove this `mut`
    |
    = note: `#[warn(unused_mut)]` on by default

warning: variable does not need to be mutable
   --> src/main.rs:334:9
    |
334 |     let mut aa = 1;
    |         ----^^
    |         |
    |         help: remove this `mut`

warning: function is never used: `next_permutation`
   --> src/main.rs:145:4
    |
145 | fn next_permutation(a: &mut Vec<usize>) -> bool {
    |    ^^^^^^^^^^^^^^^^
    |
    = note: `#[warn(dead_code)]` on by default

warning: function is never used: `find`
   --> src/main.rs:178:4
    |
178 | fn find(a: &mut Vec<i64>, p: i64) -> bool {
    |    ^^^^

warning: function is never used: `solve`
   --> src/main.rs:203:4
    |
203 | fn solve(idx: usize, last: i64, n: usize, p: i64, a: &mut Vec<i64>) -> i64 {
    |    ^^^^^

warning: function is never used: `experiment`
   --> src/main.rs:244:4
    |
244 | fn experiment(n: usize, p: i64) -> i64 {
    |    ^^^^^^^^^^

warning: function is never used: `ex3`
   --> src/main.rs:249:4
    |
249 | fn ex3(n1: usize, n2: usize, n3: usize, n4: usize, n: usize, p: usize) -> bool {
    |    ^^^

warning: function is never used: `ex2`
   --> src/main.rs:286:4
    |
286 | fn ex2(n: usize, p: usize) {
    |    ^^^

warning: comparison is useless due to type limits
   --> src/main.rs:151:11
    |
151 |     while x >= 0 && a[x] >= a[x+1] {
    |          ...

ジャッジ結果