AtCoder Beginner Contest 359 D

https://atcoder.jp/contests/abc359/tasks/abc359_d

単純なDPですね。K文字でAを0、Bを1と考えれば、0～2^K-1までの整数が状態となります。

// Avoid K Palindrome
#![allow(non_snake_case)]


//////////////////// library ////////////////////

fn read<T: std::str::FromStr>() -> T {
    let mut line = String::new();
    std::io::stdin().read_line(&mut line).ok();
    line.trim().parse().ok().unwrap()
}

fn read_vec<T: std::str::FromStr>() -> Vec<T> {
    read::<String>().split_whitespace()
            .map(|e| e.parse().ok().unwrap()).collect()
}


//////////////////// process ////////////////////

fn read_input() -> (usize, String) {
    let v: Vec<usize> = read_vec();
    let K = v[1];
    let S: String = read();
    (K, S)
}

const D: i64 = 998244353;

type DP = Vec<i64>;

fn is_palindromic(n: usize, K: usize) -> bool {
    for i in 0..K/2 {
        if ((n >> i) & 1) != ((n >> (K-i-1)) & 1) {
            return false;
        }
    }
    true
}

fn initialize_DP(K: usize, S: &String) -> DP {
    let mut v: Vec<usize> = vec![0];
    for i in 0..K {
        let c = S.as_bytes()[i] as char;
        let mut w = vec![];
        for s in v.into_iter() {
            if c != 'B' {
                w.push(s << 1)
            }
            if c != 'A' {
                w.push((s << 1) | 1)
            }
        }
        v = w
    }
    
    let mut dp: DP = vec![0; 1<<K];
    for s in v.into_iter() {
        if !is_palindromic(s, K) {
            dp[s] = 1
        }
    }
    dp
}

fn update_DP(dp: DP, c: char, K: usize) -> DP {
    let mask = (1 << K) - 1;
    let mut new_dp: DP = vec![0; 1<<K];
    for s0 in 0..1<<K {
        if dp[s0] == 0 {
            continue
        }
        if c != 'B' {
            let s = (s0 << 1) & mask;
            if !is_palindromic(s, K) {
                new_dp[s] = (new_dp[s] + dp[s0]) % D
            }
        }
        if c != 'A' {
            let s = ((s0 << 1) | 1) & mask;
            if !is_palindromic(s, K) {
                new_dp[s] = (new_dp[s] + dp[s0]) % D
            }
        }
    }
    new_dp
}

fn F(K: usize, S: String) -> i64 {
    let mut dp = initialize_DP(K, &S);
    for i in K..S.len() {
        dp = update_DP(dp, S.as_bytes()[i] as char, K)
    }
    dp.into_iter().sum::<i64>() % D
}

fn main() {
    let (K, S) = read_input();
    println!("{}", F(K, S))
}