Problem:
H: Revenge of UMG

問題

'U', 'M', 'G' の三種類の文字からなる文字列 T に対し、1, 2, ..., |T| 文字目をそれぞれ T_1, T_2, ..., T_{|T|} と表すことにしたとき、以下の条件を満たす (i, j, k) の組の個数を、文字列 T の “UMG 数”と呼ぶことにします:

1 \leq i < j < k \leq |T|

j - i = k - j

T_i = 'U', T_j = 'M', T_k = 'G'

さて、'U', 'M', 'G', '?' の四種類の文字からなる文字列 S が与えられます。S の '?' をそれぞれ 'U', 'M', 'G' のいずれかに置き換えてできる文字列は、'?' の個数を N として 3^{N} 通り考えられますが、それぞれの文字列の UMG 数について総和をとったものを 998244353 で割ったあまりを求めてください。

入力形式

S

制約

3 \leq |S| \leq 2 \times 10^{5}

S は 'U', 'M', 'G', '?' の四種類の文字からなる文字列である。

出力形式

答えを表す整数を一行に出力してください。998244353 で割ったあまりを出力することに注意してください。

入力例 1

?MG?

出力例 1

3

最初の '?' が 'U' でない場合には UMG 数は 0 になります。最初の '?' を 'U' としたとき、

UMGU の UMG 数は 1

UMGM の UMG 数は 1

UMGG の UMG 数は 1

となり、合計値は 3 になります。

入力例 2

UUMMGGGUMG

出力例 2

4

入力例 3

?????G????U???M??????G?????M????G???U??????M???G??

出力例 3

648330088

998244353 で割ったあまりを答えてください。

Solution in Rust (judge verdict: Accepted)

const MOD: u32 = 998_244_353;
#[derive(Clone, Copy)]
struct ModInt(u32);
impl std::ops::Add for ModInt {
    type Output = ModInt;
    fn add(self, rhs: ModInt) -> ModInt {
        let mut d = self.0 + rhs.0;
        if d >= MOD {
            d -= MOD;
        }
        ModInt(d)
    }
}

impl std::ops::Mul for ModInt {
    type Output = ModInt;
    fn mul(self, rhs: ModInt) -> ModInt {
        ModInt((self.0 as u64 * rhs.0 as u64 % MOD as u64) as u32)
    }
}

impl std::ops::Neg for ModInt {
    type Output = ModInt;
    fn neg(self) -> ModInt {
        if self.0 == 0 {self} else {ModInt(MOD - self.0)}
    }
}

impl std::ops::AddAssign for ModInt {
    fn add_assign(&mut self, rhs: ModInt) {
        *self = *self + rhs;
    }
}

impl std::ops::MulAssign for ModInt {
    fn mul_assign(&mut self, rhs: ModInt) {
        *self = *self * rhs;
    }
}

#[allow(dead_code)]
impl ModInt {
    pub fn new(v: u32) -> ModInt {
        ModInt(v % MOD)
    }
    pub fn zero() -> ModInt {
        ModInt(0)
    }
    pub fn one() -> ModInt {
        ModInt(1)
    }
    pub fn pow(self, mut n: u32) -> ModInt {
        let mut t = ModInt::one();
        let mut s = self;
        while n > 0 {
            if n & 1 == 1 {
                t = t * s;
            }
            s = s * s;
            n >>= 1;
        }
        t
    }
    pub fn inv(self) -> ModInt {
        self.pow(MOD - 2)
    }
}

type M = ModInt;

fn fft(zeta: M, n: usize, f: &mut [M], g: &mut [M]) {
    if n == 1 {
        g[0] = f[0];
        return;
    }
    let m = n / 2;
    for i in 0..m {
        g[i] = f[2 * i];
        g[m + i] = f[2 * i + 1];
    }
    let zeta2 = zeta * zeta;
    fft(zeta2, m, &mut g[0..m], &mut f[0..m]);
    fft(zeta2, m, &mut g[m..n], &mut f[m..n]);
    let mut q = M::one();
    for (g, (a, b)) in g.iter_mut().zip(f[0..m].iter().zip(f[m..n].iter()).cycle()) {
        *g = *a + q * *b;
        q = q * zeta;
    }
}

fn multiply(a: &[M], b: &[M]) -> Vec<M> {
    if std::cmp::min(a.len(), b.len()) < 32 {
        let mut c = vec![M::zero(); a.len() + b.len() - 1];
        for (i, a) in a.iter().enumerate() {
            for (c, b) in c[i..].iter_mut().zip(b.iter()) {
                *c += *a * *b;
            }
        }
        return c;
    }
    let n = a.len() + b.len();
    let mut k = 1;
    while k < n - 1 {
        k *= 2;
    }
    let mut f = vec![M::zero(); k];
    let mut g = vec![M::zero(); k];
    let mut x = Vec::with_capacity(k);
    x.resize(k, M::zero());
    for (x, a) in x.iter_mut().zip(a.iter()) {
        *x = *a;
    }
    let zeta = M::new(15311432).pow(((1 << 23) / k) as u32);
    fft(zeta, k, &mut x, &mut f);
    x.clear();
    x.resize(k, M::zero());
    for (x, b) in x.iter_mut().zip(b.iter()) {
        *x = *b;
    }
    fft(zeta, k, &mut x, &mut g);
    for (x, (f, g)) in x.iter_mut().zip(f.iter().zip(g.iter())) {
        *x = *f * *g;
    }
    let izeta = zeta.pow(MOD - 2);
    fft(izeta, k, &mut x, &mut f);
    let mut ans = vec![M::zero(); n - 1];
    let inv = M::new(k as u32).pow(MOD - 2);
    for i in 0..(n - 1) {
        ans[i] = f[i] * inv;
    }
    ans
}

fn calc(a: &[ModInt], b: &[ModInt], c: &mut [ModInt]) {
    if a.len() <= 2 {
        return;
    }
    let n = a.len();
    let m = n / 2;
    calc(&a[0..m], &b[0..m], &mut c[0..(2 * m)]);
    calc(&a[m..n], &b[m..n], &mut c[(2 * m)..(2 * n)]);
    let x = multiply(&a[0..m], &b[m..n]);
    for (c, x) in c[m..].iter_mut().zip(x.iter()) {
        *c += *x;
    }
}

fn run() {
    let mut s = String::new();
    std::io::stdin().read_line(&mut s).unwrap();
    let s: Vec<char> = s.trim().chars().collect();
    let n = s.len();
    let i3 = ModInt(3).inv();
    let mut a = vec![ModInt::zero(); n];
    let mut b = vec![ModInt::zero(); n];
    let mut cnt = 0;
    for i in 0..n {
        match s[i] {
            'U' => a[i] = ModInt::one(),
            'G' => b[i] = ModInt::one(),
            '?' => {
                cnt += 1;
                a[i] = i3;
                b[i] = i3;
            },
            _ => (),
        }
    }
    let mut c = vec![ModInt::zero(); 2 * n];
    calc(&a, &b, &mut c);
    let mut ans = ModInt::zero();
    for i in 0..n {
        if s[i] == 'M' {
            ans += c[2 * i];
        } else if s[i] == '?' {
            ans += c[2 * i] * i3;
        }
    }
    ans *= ModInt(3).pow(cnt);
    println!("{}", ans.0);
}

fn main() {
    run();
}