AtCoder Beginner Contest 331 復習

Rustで競プロするのに慣れたりライブラリを整備したりするのも兼ねて、ABC331を解き直します。
コンテスト時はA～Eの5完でした。

自作ライブラリを使用し、ローカルで事前に展開する形で提出しています。展開前のコードを記述します。
AtCoderで使えるクレートに含まれない関数などを使っている可能性があること、破壊的変更をこの記事の投稿後にしている可能性があることにご注意ください。
github.com

A問題

atcoder.jp
月に対して日数が固定で助かりました。まあA問題ですからね。
コンテスト中の提出からは変数名だけを変えています。この問題ではこの実装が最もシンプルだと思います。
atcoder.jp

#[allow(unused_attributes)] #[macro_use] #[allow(unused_imports)] use not_leonian_ac_lib::*;

fn main() {
    /// 標準入力のマクロ（インタラクティブ問題ならば中のマクロを変更）
    macro_rules! input {
        ($($tt:tt)*) => {
            proconio::input!($($tt)*);
            // proconio::input_interactive!($($tt)*);
        };
    }

    input!(mnum:usize,dnum:usize,mut y:usize,mut m:usize,mut d:usize);
    d+=1;
    if d>dnum {
        m+=1;
        d=1;
    }
    if m>mnum {
        y+=1;
        m=1;
    }
    outputln!(y,m,d);
}

B問題

atcoder.jp
本番なぜか全探索が思い浮かばず、結局ナップサックDPを書いてしまいました。
B問題で21:10を越えたのは久しぶりですね。
全探索の際の個数をnに合わせると+1などによる混乱は起こりうるので、本番ならとりあえずで全部100ずつにするのもありでしょうね。
atcoder.jp

#[allow(unused_attributes)] #[macro_use] #[allow(unused_imports)] use not_leonian_ac_lib::*;

fn main() {
    /// 標準入力のマクロ（インタラクティブ問題ならば中のマクロを変更）
    macro_rules! input {
        ($($tt:tt)*) => {
            proconio::input!($($tt)*);
            // proconio::input_interactive!($($tt)*);
        };
    }

    input!(n:usize,s:usize,m:usize,l:usize);
    let mut ans=E[18];
    for i in 0..=n/6+1 {
        for j in 0..=n/8+1 {
            for k in 0..=n/12+1 {
                if 6*i+8*j+12*k>=n {
                    ans.chmin(s*i+m*j+l*k);
                }
            }
        }
    }
    outputln!(ans);
}

C問題

atcoder.jp
本番は累積和と二分探索を行いました。もっと綺麗に実装できるかもしれないとは思ってましたが、実際公式解説ではなかったですね。
元の添字とともにソートをする典型を行い、求めるのはA_iより大きな要素全ての和なので、新たな数が出てきたら今時点での総和を代入し、前と同じ数ならば前と同じ総和を代入すればよいです。
降順ソートはとりあえずライブラリに入れないでおきます。そろそろ文字数を減らすだけの、典型アルゴリズムでない関数はライブラリから省きたいくらいの行数になってしまっています。
atcoder.jp

#[allow(unused_attributes)] #[macro_use] #[allow(unused_imports)] use not_leonian_ac_lib::*;

fn main() {
    /// 標準入力のマクロ（インタラクティブ問題ならば中のマクロを変更）
    macro_rules! input {
        ($($tt:tt)*) => {
            proconio::input!($($tt)*);
            // proconio::input_interactive!($($tt)*);
        };
    }

    input!(n:usize,a:[usize;n]);
    let mut asort=vec_range(0, n, |i| (a[i],i));
    asort.sort();
    asort.reverse();
    let mut b=vec![0;n];
    let mut sum=asort[0].0;
    for i in 1..n {
        if asort[i].0!=asort[i-1].0 {
            b[asort[i].1]=sum;
        } else {
            b[asort[i].1]=b[asort[i-1].1];
        }
        sum+=asort[i].0;
    }
    b.outputln();
}

D問題

atcoder.jp
本番は本当に恥ずかしいムーブをしてしまいました。9分割して40分くらいかけてしまいました。公式解説の方法を思いついていればレートの微減はなかったかもしれません。
2次元累積和で(0,0)から(i-1,j-1)までの長方形領域の黒いマスの個数を求めつつ、(a,b)から(c,d)までの長方形領域の黒いマスの個数は同じく2次元累積和の考え方で足し引きすればよいです。
2次元累積和で(0,0)から(i-1,j-1)までの長方形領域の黒いマスの個数を求める際は、巡回している部分と余りで4分割すればよいです。
私は2次元累積和をライブラリ化しているので、その関数を使ったクロージャを書きました。
atcoder.jp

#[allow(unused_attributes)] #[macro_use] #[allow(unused_imports)] use not_leonian_ac_lib::*;
use proconio::marker::Chars;

fn main() {
    /// 標準入力のマクロ（インタラクティブ問題ならば中のマクロを変更）
    macro_rules! input {
        ($($tt:tt)*) => {
            proconio::input!($($tt)*);
            // proconio::input_interactive!($($tt)*);
        };
    }

    input!(n:usize,q:usize,p:[Chars;n],abcd:[(usize,usize,usize,usize);q]);
    let color=vec_range(0, n, |i| vec_range(0, n, |j| (p[i][j]=='B') as usize));
    let sums=color.construct_2d_prefix_sum();
    let prefix_sum=|a:usize,b:usize| {
        let mut ret=0;
        ret+=sums.calculate_2d_partial_sum(0, 0, n, n)*(a/n)*(b/n);
        ret+=sums.calculate_2d_partial_sum(0, 0, a%n, n)*(b/n);
        ret+=sums.calculate_2d_partial_sum(0, 0, n, b%n)*(a/n);
        ret+=sums.calculate_2d_partial_sum(0, 0, a%n, b%n);
        ret
    };
    for (a,b,c,d) in abcd {
        let mut ans=0;
        ans+=prefix_sum(c+1,d+1);
        ans+=prefix_sum(a,b);
        ans-=prefix_sum(a,d+1);
        ans-=prefix_sum(c+1,b);
        outputln!(ans);
    }
}

E問題

atcoder.jp
本番はevimaさんの解説と同じ方法で解きました。この解法が好きなのでこの記事でもそれで書きます。
主菜を全探索し、メニューにあるものが見つかるまで高い副菜を全探索します。
探索するのは、各主菜について組み合わせがメニューにある最も高い副菜と、それより高く組み合わせがメニューにない副菜のみなので、最大でも(N+L)個です。

コードについて、コンテスト時の提出からは変数名やfor文の宣言のみ少し変えました。
atcoder.jp

use std::collections::BTreeSet;

#[allow(unused_attributes)] #[macro_use] #[allow(unused_imports)] use not_leonian_ac_lib::*;
use proconio::marker::Usize1;

fn main() {
    /// 標準入力のマクロ（インタラクティブ問題ならば中のマクロを変更）
    macro_rules! input {
        ($($tt:tt)*) => {
            proconio::input!($($tt)*);
            // proconio::input_interactive!($($tt)*);
        };
    }

    input!(n:usize,m:usize,l:usize,a:[usize;n],b:[usize;m],cd:[(Usize1,Usize1);l]);
    let mut mp=BTreeSet::<(usize,usize)>::new();
    for cd in cd {
        mp.insert(cd);
    }
    let mut bnum=vec_range(0, m, |i| (b[i],i));
    bnum.sort();
    let mut ans=0;
    for i in 0..n {
        for &(b,j) in bnum.iter().rev() {
            if !mp.contains(&(i,j)) {
                ans.chmax(a[i]+b);
                break;
            }
        }
    }
    outputln!(ans);
}

F問題

atcoder.jp
今まで見たローリングハッシュを使える問題は全部ローリングハッシュが別解だったので勉強を避けていたのが仇になってしまいました。
累積和とmod逆元を使ったローリングハッシュの実装法があることを知っていれば、セグ木に乗せる発想に本番たどり着けていたかもしれません。
コンテスト中はManacherのアルゴリズムを使って回文の半径を求め、その値を変化させるという発想からほぼ抜け出せませんでした。
ローリングハッシュはそのままではセグ木に乗りませんが、基数の長さぶんの冪を一緒に乗せることでモノイドになります。
この前のこれと同じ発想であるだけに思いつけなかったのは悔しいですね。
atcoder.jp
この機会にローリングハッシュとセグ木に乗せる構造をライブラリに入れました。
（入水からちょっと経ったあたりで大負けしてトラウマの）ABC284のF問題で先に普通のローリングハッシュを試しましたが、だいぶバグらせました。
この問題はクエリ平方分割対策でNの制約が厳しいのもありますが、最初基数の個数を5個で固定していたら2.5s超えて危なかったので、基数の個数は指定できるようにしました。
atcoder.jp

use ac_library::Segtree;
#[allow(unused_attributes)] #[macro_use] #[allow(unused_imports)] use not_leonian_ac_lib::*;
use proconio::marker::{Chars, Usize1};

fn main() {
    /// 標準入力のマクロ（インタラクティブ問題ならば中のマクロを変更）
    macro_rules! input {
        ($($tt:tt)*) => {
            proconio::input!($($tt)*);
            // proconio::input_interactive!($($tt)*);
        };
    }

    input!(n:usize,q:usize,s:Chars);
    let bases=RollingHashBases::<3>::gen(n, 26).bases();
    let mut segs=Segtree::<DynamicRollingHash<3>>::new(n);
    let mut revs=Segtree::<DynamicRollingHash<3>>::new(n);
    for i in 0..n {
        segs.set(i, ([s[i].from_lowercase()+1;3], bases));
        revs.set(n-1-i, ([s[i].from_lowercase()+1;3], bases))
    }
    for _ in 0..q {
        input!(k:usize);
        if k==1 {
            input!(x:Usize1,c:char);
            segs.set(x, ([c.from_lowercase()+1;3], bases));
            revs.set(n-1-x, ([c.from_lowercase()+1;3], bases))
        } else {
            input!(l:Usize1,r:Usize1);
            output_yes_or_no(segs.prod(l..=r)==revs.prod(n-1-r..=n-1-l));
        }
    }
}

G問題

atcoder.jp
明らかに自分の実力よりは上のレベルの問題ですが、実装は大変ではないという噂を聞いたことと、下の別解に感動したことから取り組んでみたいと思います。
（最初、現在は修正されている記事の間違いに気付かず拡散してしまい、申し訳ないです。）
rogi52.hatenablog.jp
not-leonian.hatenablog.com
簡単に言えば、等確率でないガチャのコンプまでの回数の期待値を求める問題です。
和事象・積事象と確率の加法定理は上の私の記事の一般化した包除原理の種ですから、包除原理を使えます。したがって、以下のような変形ができます。
A_iを、iがノートに書かれているという事象とします。また、p_k(A)を、k回目の試行の後に事象Aが起こる確率とします。
\displaystyle \begin{align} E{\left[\bigcap_{i\in\{1,2,3,\ldots,M\}}A_i\right]}&=\sum_{k=1}^{\infty} kp_k{\left({\left[\bigcap_{i\in\{1,2,3,\ldots,M\}}A_i\right]}^c\right)}\\&=\sum_{k=1}^{\infty} kp_k{\left(\bigcup_{i\in\{1,2,3,\ldots,M\}}{A_i}^c\right)}\\&=-\sum_{k=1}^{\infty} \left[k\sum_{\emptyset\varsubsetneq S\subseteq\{1,2,3,\ldots,M\}}(-1)^{|S|}p_k{\left(\bigcap_{i\in S}{A_i}^c\right)}\right]\\&=-\sum_{\emptyset\varsubsetneq S\subseteq\{1,2,3,\ldots,M\}}(-1)^{|S|}\sum_{k=1}^{\infty} kp_k{\left(\bigcap_{i\in S}{A_i}^c\right)}\\&=-\sum_{\emptyset\varsubsetneq S\subseteq\{1,2,3,\ldots,M\}}(-1)^{|S|}\sum_{k=1}^{\infty} kp_k{\left({\left[\bigcup_{i\in S}A_i\right]}^c\right)}\\&=-\sum_{\emptyset\varsubsetneq S\subseteq\{1,2,3,\ldots,M\}}(-1)^{|S|}E{\left[\bigcup_{i\in S}A_i\right]} \end{align}
ただし、途中に以下の問題の解説のような式変形を用いました。また、途中の級数と総和の順序交換には、各級数の収束と級数の線形性を用いています。
atcoder.jp

ここで右辺に現れる期待値は、Sに含まれるいずれかの数がノートに書かれるまでの試行回数の期待値です。
そして、Sに含まれるいずれかの数がノートに書かれる確率は容易に\sum_{i\in S} C_i/Nとわかります。
そして、試行回数の期待値はその逆数です。有名な定理でもありますし、最終的な式の直前の式に出てくる確率の級数からも容易に求まります。
あとは主客転倒と形式的冪級数の展開を行えば、FFTマージテクが使える形に変形できます。
\displaystyle \begin{align} E{\left[\bigcap_{i\in\{1,2,3,\ldots,M\}}A_i\right]}&=-\sum_{\emptyset\varsubsetneq S\subseteq\{1,2,3,\ldots,M\}}(-1)^{|S|}E{\left[\bigcup_{i\in S}A_i\right]}\\&=-\sum_{\emptyset\varsubsetneq S\subseteq\{1,2,3,\ldots,M\}}(-1)^{|S|}\frac{N}{\sum_{i\in S} C_i}\\&=-N\sum_{k=1}^N \sum_{\sum_{i\in S} C_i=k}(-1)^{|S|}\frac{1}{k}\\&=-N\sum_{k=1}^N \frac{1}{k}\sum_{\sum_{i\in S} C_i=k}(-1)^{|S|}\\&=-N\sum_{k=1}^N \frac{1}{k}\left\{[X^k]\prod_{i=1}^M\left(1-X^{C_i}\right)\right\} \end{align}

FPSライブラリは元々作っていたので、FFTマージテクを実装しました。
意外と簡単に実装できたのですが、四則演算を両方の次数が等しいときにしか行えないようにしていたので、それを直すのに一苦労でした。
atcoder.jp

use ac_library::ModInt998244353;
#[allow(unused_attributes)] #[macro_use] #[allow(unused_imports)] use not_leonian_ac_lib::*;

fn main() {
    /// 標準入力のマクロ（インタラクティブ問題ならば中のマクロを変更）
    macro_rules! input {
        ($($tt:tt)*) => {
            proconio::input!($($tt)*);
            // proconio::input_interactive!($($tt)*);
        };
    }

    input!(n:usize,m:usize,c:[usize;m]);
    let mut ans=ModInt998244353::new(0);
    let mut fps=vec_range(0, m, |i| vec_range(0, c[i]+1, |j| ModInt998244353::new(
        if j==0 {
            1
        } else if j==c[i] {
            -1
        } else {
            0
        }
    )));
    FPS::fps_prod_merge(&mut fps);
    for i in 1..=n {
        let mut inv=ModInt998244353::new(1);
        inv/=i;
        ans+=inv*fps[0][i];
    }
    ans*=n;
    ans*=-1;
    outputln!(ans);
}

また、exp-log典型を用いることもできます。
\displaystyle \begin{align} \prod_{i=1}^M\left(1-X^{C_i}\right)&=\exp{\left(\sum_{i=1}^M\log{\left(1-X^{C_i}\right)}\right)}\\&=\exp{\left(\sum_{k=1}^N \left[\Big|\{C_i|C_i=k\}\Big|\cdot\log{\left(1-X^k\right)}\right]\right)}\\&=\exp{\left(-\sum_{k=1}^N \left[\Big|\{C_i|C_i=k\}\Big|\sum_{j=1}^{\infty}\frac{X^{jk}}{j}\right]\right)} \end{align}
より、\Theta(N\log N)で計算することができます。
ただしあくまでオーダーが小さくなるだけで、expが重いぶん今回の制約での実行時間は長くなります。

FPSライブラリを作っていたはいいものの、実際本番で通せたことがないのでかなりひどい実装になっていました。
四則演算を両方の次数が等しくなくても使えるようにしたぶん、expやlogを直すのにひどくバグらせてしまいました。
atcoder.jp

use ac_library::ModInt998244353;
#[allow(unused_attributes)] #[macro_use] #[allow(unused_imports)] use not_leonian_ac_lib::*;

fn main() {
    /// 標準入力のマクロ（インタラクティブ問題ならば中のマクロを変更）
    macro_rules! input {
        ($($tt:tt)*) => {
            proconio::input!($($tt)*);
            // proconio::input_interactive!($($tt)*);
        };
    }

    input!(n:usize,m:usize,c:[usize;m]);
    let mut ans=ModInt998244353::new(0);
    let mut cnts=vec![0;n+1];
    for c in c {
        cnts[c]+=1;
    }
    let mut invs=vec![ModInt998244353::new(1);n+1];
    for i in 1..=n {
        invs[i]/=i;
    }
    let mut log=vec![ModInt998244353::new(0);n+1];
    for i in 1..=n {
        for j in 1..=n/i {
            log[i*j]-=invs[j]*cnts[i];
        }
    }
    let fps=log.fps_exp(n);
    for i in 1..=n {
        ans+=invs[i]*fps[i];
    }
    ans*=n;
    ans*=-1;
    outputln!(ans);
}