AtCoder Beginner Contest 322 復習 (Rust)

Rustで競プロするのに慣れたりライブラリを整備したりするのも兼ねて、ABC322を解き直します。

A問題

atcoder.jp
本番はスライスで実装しました。これはもう書き直すまでもないでしょう。
atcoder.jp

fn main() {
    macro_rules! input {
        ($($tt:tt)*) => {
            proconio::input!($($tt)*);
        };
    }

    #[allow(unused_macros)]
    macro_rules! outputln {
        ($var:expr) => {
            println!("{}",$var);
        };
        ($var:expr,$($vars:expr),+) => {
            print!("{} ",$var);
            outputln!($($vars),+);
        };
    }

    input!(n:usize,s:String);
    for i in 0..n-2 {
        if &s[i..i+3]=="ABC" {
            outputln!(i+1);
            return;
        }
    }
    outputln!(-1);
}

B問題

atcoder.jp
本番はA問題同様スライスで実装しましたが、後でstarts_with関数とends_with関数を知りました。
可読性ならmatchですが記述量ならas usize。
atcoder.jp

fn main() {
    macro_rules! input {
        ($($tt:tt)*) => {
            proconio::input!($($tt)*);
        };
    }

    #[allow(unused_macros)]
    macro_rules! outputln {
        ($var:expr) => {
            println!("{}",$var);
        };
        ($var:expr,$($vars:expr),+) => {
            print!("{} ",$var);
            outputln!($($vars),+);
        };
    }

    input!(_:usize,_:usize,s:String,t:String);
    outputln!(3-t.starts_with(&s) as usize*2-t.ends_with(&s) as usize);
}

C問題

atcoder.jp
公式解説とは違う方法で解きました。i日目以降に花火が初めて上がるのがA_m日目の場合、(i+1)日目以降に花火が初めて上がるのは必ずA_m日目またはA_{m+1}日目であること、N日目に必ず花火が上がることを利用して、前から花火が初めて上がる日を更新していきました。
公式解説とは違う方法ですが個人的には自分の解法のほうがしっくりくるので、これも書き直さなくてよいでしょう。
atcoder.jp

fn main() {
    macro_rules! input {
        ($($tt:tt)*) => {
            proconio::input!($($tt)*);
        };
    }

    #[allow(unused_macros)]
    macro_rules! outputln {
        ($var:expr) => {
            println!("{}",$var);
        };
        ($var:expr,$($vars:expr),+) => {
            print!("{} ",$var);
            outputln!($($vars),+);
        };
    }

    input!(n:usize,m:usize,a:[usize;m]);
    let mut cur=0;
    for i in 1..=n {
        if i>a[cur] {
            cur+=1;
        }
        outputln!(a[cur]-i);
    }
}

D問題

atcoder.jp
流石にこれ書き直しは勘弁させて…という感想ですが、回転を4通り書くのではなく90°回転だけを書くことで実装ミスを減らすのはなるほどと思ったのでそこだけ書き直し。
Rustの速度でごり押してしまいました。
atcoder.jp

use proconio::marker::Chars;

fn main() {
    macro_rules! input {
        ($($tt:tt)*) => {
            proconio::input!($($tt)*);
        };
    }

    #[allow(unused_macros)]
    macro_rules! outputln {
        ($var:expr) => {
            println!("{}",$var);
        };
        ($var:expr,$($vars:expr),+) => {
            print!("{} ",$var);
            outputln!($($vars),+);
        };
    }

    input!(p1:[Chars;4],p2:[Chars;4],p3:[Chars;4]);
    let mut p2s=vec![vec![vec!['.';4];4];4];
    let mut p3s=vec![vec![vec!['.';4];4];4];
    for i in 0..4 {
        for j in 0..4 {
            p2s[0][i][j]=p2[i][j];
            p3s[0][i][j]=p3[i][j];
        }
    }
    for k in 1..4 {
        for i in 0..4 {
            for j in 0..4 {
                p2s[k][j][3-i]=p2s[k-1][i][j];
                p3s[k][j][3-i]=p3s[k-1][i][j];
            }
        }
    }
    for r2 in 0..4 {
        for r3 in 0..4 {
            for i1 in 0..8 {
                for i2 in 0..8 {
                    for i3 in 0..8 {
                        for j1 in 0..8 {
                            for j2 in 0..8 {
                                for j3 in 0..8 {
                                    let mut ok=true;
                                    let mut g=vec![vec!['x';12];12];
                                    for i in 4..8 {
                                        for j in 4..8 {
                                            g[i][j]='.';
                                        }
                                    }
                                    for i in 0..4 {
                                        for j in 0..4 {
                                            if p1[i][j]!='.' && g[i+i1][j+j1]!='.' {
                                                ok=false;
                                            }
                                            g[i+i1][j+j1]=p1[i][j];
                                        }
                                    }
                                    for i in 0..4 {
                                        for j in 0..4 {
                                            if p2s[r2][i][j]!='.' && g[i+i2][j+j2]!='.' {
                                                ok=false;
                                            }
                                            if p2s[r2][i][j]!='.' {
                                                g[i+i2][j+j2]=p2s[r2][i][j];
                                            }
                                        }
                                    }
                                    for i in 0..4 {
                                        for j in 0..4 {
                                            if p3s[r3][i][j]!='.' && g[i+i3][j+j3]!='.' {
                                                ok=false;
                                            }
                                            if p3s[r3][i][j]!='.' {
                                                g[i+i3][j+j3]=p3s[r3][i][j];
                                            }
                                        }
                                    }
                                    for i in 4..8 {
                                        for j in 4..8 {
                                            if g[i][j]=='.' {
                                                ok=false;
                                            }
                                        }
                                    }
                                    if ok {
                                        outputln!("Yes");
                                        return;
                                    }
                                }
                            }
                        }
                    }
                }
            }
        }
    }
    outputln!("No");
}

E問題

atcoder.jp
本番はK=5で固定して、最後に答えを参照する添字のみを変える脳筋戦法で通しました。
公式解説の方法で書いてみましたがミスの危険性はむしろ増えそうな…
Rustの速度を信じてBTreeMapが良いのでしょうか。
あと、RustでのDPテーブルの更新のやりにくさは、当分の間1次元なら自作関数を使用、2次元以上ならsplit_at_mut関数を使用で対応してみます。
atcoder.jp

fn main() {
    macro_rules! input {
        ($($tt:tt)*) => {
            proconio::input!($($tt)*);
        };
    }

    #[allow(unused_macros)]
    macro_rules! outputln {
        ($var:expr) => {
            println!("{}",$var);
        };
        ($var:expr,$($vars:expr),+) => {
            print!("{} ",$var);
            outputln!($($vars),+);
        };
    }

    input!(n:usize,k:usize,p:usize);
    let mut c=vec![0;n];
    let mut a=vec![vec![0;k];n];
    for i in 0..n {
        input!(ci:usize);
        c[i]=ci;
        input!(ai:[usize;k]);
        a[i]=ai;
    }
    let pmax=(p+1).pow(k as u32);
    let mut dp=vec![vec![usize::MAX;pmax];n+1];
    dp[0][0]=0;
    for i in 0..n {
        for ps in 0..pmax {
            let (prev,next)=dp.split_at_mut(i+1);
            next[0][ps].chmin(prev[i][ps]);
            let pvec=vec_range(0, k, |j| {
                ps/((p+1).pow(j as u32))%(p+1)
            });
            let mut newps=0;
            for j in 0..k {
                newps+=min(pvec[j]+a[i][j],p)*((p+1).pow(j as u32));
            }
            next[0][newps].chmin(prev[i][ps].saturating_add(c[i]));
        }
    }
    dp[n][pmax-1].outputifexists(usize::MAX);
}

trait Outputifexists {
    fn outputifexists(self, max: Self);
}

impl Outputifexists for usize {
    fn outputifexists(self, max: Self) {
        if self<max {
            println!("{}",self);
        } else {
            println!("-1");
        }
    }
}

#[allow(dead_code)]
fn vec_range<N,F,T>(begin: N, end: N, func: F) -> Vec<T> where std::ops::Range<N>: Iterator, F: Fn(<std::ops::Range<N> as Iterator>::Item) -> T {
    return (begin..end).map(|i| func(i)).collect::<Vec::<T>>();
}

#[allow(dead_code)]
fn min<T>(left: T, right: T) -> T where T: std::cmp::PartialOrd {
    return if left<right {
        left
    } else {
        right
    };
}

trait Chminmax {
    fn chmin(&mut self, challenger: Self);
}

impl<T> Chminmax for T where T: Copy + std::cmp::PartialOrd {
    fn chmin(&mut self, challenger: Self) {
        if challenger<*self {
            *self=challenger;
        }
    }
}

F問題

atcoder.jp
見た目は暖色diffの問題ですが、中国では典型だったようです。
遅延セグ木を使うことと、答え以外に両端の1の長さをもつこと、1だけでなく0についての情報ももつことまでは時間内に考察できましたが、モノイドにするためにもう1つ区間長といった情報が必要なことがわかりませんでした。
Rustに移行してから遅延セグ木を使ったことがなかったので、どちらにしても実装は間に合わなかった気がします。
結論から言うと、遅延セグ木に乗せる情報は「区間内の1のみからなる最長の区間の長さ」「区間の左端から1が続く長さ」「区間の右端から1が続く長さ」「区間内の0のみからなる最長の区間の長さ」「区間の左端から0が続く長さ」「区間の右端から0が続く長さ」「区間の長さ」でした。
区間Iに対応するそれぞれの情報をm_1(I),\ l_1(I),\ r_1(I),\ m_0(I),\ l_0(I),\ r_0(I),\ s(I)とおきます。
I=I_1\sqcup I_2 (ただし\forall n_1\in I_1.\forall n_2\in I_2.n_1\lt n_2)について、I_1,\ I_2の値からIの値を求めることを考えます。
I内で1または0のみからなる最長の区間は、I_1内の最長の区間またはI_2内の最長の区間またはI_1とI_2をまたがっている区間なので、それぞれm_1(I)=\max\{m_1(I_1),\ m_1(I_2),\ r_1(I_1)+l_1(I_2)\},\ m_0(I)=\max\{m_0(I_1),\ m_0(I_2),\ r_0(I_1)+l_0(I_2)\}となります。
Iの端から1または0が続く長さは、その端側の区間の全てで続いているかによって場合分けする必要があります。
具体的には、\displaystyle l_1(I)=\begin{cases}l_1(I_1)&(l_1(I_1)\lt s(I_1))\\s(I_1)+l_1(I_2)&(l_1(I_1)=s(I_1))\end{cases},\ r_1(I)=\begin{cases}r_1(I_2)&(r_1(I_2)\lt s(I_2))\\s(I_2)+r_1(I_1)&(r_1(I_2)=s(I_2))\end{cases},\ l_0(I)=\begin{cases}l_0(I_1)&(l_0(I_1)\lt s(I_1))\\s(I_1)+l_0(I_2)&(l_0(I_1)=s(I_1))\end{cases},\ r_0(I)=\begin{cases}r_0(I_2)&(r_0(I_2)\lt s(I_2))\\s(I_2)+r_0(I_1)&(r_0(I_2)=s(I_2))\end{cases}となります。
最後に、区間の長さは明らかにs(I)=s(I_1)+s(I_2)で計算できます。
この構造について結合法則が成り立ち、区間の長さが0の場合が単位元になっていることは定義より明らかでしょう。
次に、1についての情報3つと0についての情報3つを交換して返す写像と、恒等写像が属する集合を考え、前者にtrue、後者にfalseを割り当てます。
この集合は明らかに恒等写像を要素として含みます。
また、写像の合成について明らかに閉じており、割り当てた真偽値のXORで合成写像が計算できます。
そして、モノイドの演算の形式は1についての情報と0についての情報で同じなので、足してから写像で変換するのと、足す前にそれぞれ写像で変換してから足すのとでは計算結果は変わりません。
よって、この構造を遅延セグ木に乗せることができます。
遅延セグ木の初期化では、各点に自分自身のみを含む区間についての情報を乗せていけばよいです。（本当は1点更新で初期化すると計算量のオーダーが悪くなるのですが…）
ちなみに、連続区間和maxや最大部分配列問題などと呼ばれている、区間内の連続区間和の最大値クエリに答える似た問題があるようです。


この問題も0を-\inftyで読み替えることで連続区間和maxに落とし込めるようです。（基本的に更新はできないので1についての情報と0についての情報の両方をもたないといけないのは同じです）
hotman78.hatenablog.com
連続区間和maxは、「区間内の連続区間和の最大値」「区間の左端からの連続区間和の最大値」「区間の右端からの連続区間和の最大値」「区間和」をもつことでセグ木に乗せられるようです。
求めるものが違うため微妙に乗せるものも違いますが、式の意味は似ているといえるでしょう。
実装について、ac-library-rsの遅延セグ木を使うのはRustのstruct、trait、implの仕様がわかっていないと難しいですね。（自分は自作ライブラリをなんだかんだ整備しているので苦ではありませんが…そもそも本家ACLの遅延セグ木もC++のtemplateが多少わかってないと難しいでしょうし）

use ac_library::{Monoid, MapMonoid, LazySegtree};
use proconio::marker::{Chars, Usize1};

fn main() {
    macro_rules! input {
        ($($tt:tt)*) => {
            proconio::input!($($tt)*);
        };
    }

    #[allow(unused_macros)]
    macro_rules! outputln {
        ($var:expr) => {
            println!("{}",$var);
        };
        ($var:expr,$($vars:expr),+) => {
            print!("{} ",$var);
            outputln!($($vars),+);
        };
    }

    input!(n:usize,q:usize,s:Chars);
    let mut lst=LazySegtree::<MapLensMonoid>::new(n);
    for i in 0..n {
        lst.set(i, if s[i]=='1' {
            Lengths {
                longest_1_len: 1,
                left_1_len: 1,
                right_1_len: 1,
                longest_0_len: 0,
                left_0_len: 0,
                right_0_len: 0,
                len: 1
            }
        } else {
            Lengths {
                longest_1_len: 0,
                left_1_len: 0,
                right_1_len: 0,
                longest_0_len: 1,
                left_0_len: 1,
                right_0_len: 1,
                len: 1
            }
        });
    }
    for _ in 0..q {
        input!(c:usize,l:Usize1,r:usize);
        if c==1 {
            lst.apply_range(l..r, true);
        } else {
            outputln!(lst.prod(l..r).longest_1_len);
        }
    }
}

#[derive(Clone)]
struct Lengths {
    longest_1_len: usize,
    left_1_len: usize,
    right_1_len: usize,
    longest_0_len: usize,
    left_0_len: usize,
    right_0_len: usize,
    len: usize
}

struct LensMonoid;

impl Monoid for LensMonoid {
    type S = Lengths;
    fn identity() -> Self::S {
        return Lengths {
            longest_1_len: 0,
            left_1_len: 0,
            right_1_len: 0,
            longest_0_len: 0,
            left_0_len: 0,
            right_0_len: 0,
            len: 0
        };
    }
    fn binary_operation(a: &Self::S, b: &Self::S) -> Self::S {
        return Lengths {
            longest_1_len: std::cmp::max(std::cmp::max(a.longest_1_len, b.longest_1_len), a.right_1_len+b.left_1_len),
            left_1_len: if a.len==a.left_1_len {
                a.len+b.left_1_len
            } else {
                a.left_1_len
            },
            right_1_len: if b.len==b.right_1_len {
                a.right_1_len+b.len
            } else {
                b.right_1_len
            },
            longest_0_len: std::cmp::max(std::cmp::max(a.longest_0_len, b.longest_0_len), a.right_0_len+b.left_0_len),
            left_0_len: if a.len==a.left_0_len {
                a.len+b.left_0_len
            } else {
                a.left_0_len
            },
            right_0_len: if b.len==b.right_0_len {
                a.right_0_len+b.len
            } else {
                b.right_0_len
            },
            len: a.len+b.len
        }
    }
}

struct MapLensMonoid;

impl MapMonoid for MapLensMonoid {
    type M = LensMonoid;
    type F = bool;
    fn identity_map() -> Self::F {
        return false;
    }
    fn mapping(f: &Self::F, x: &<Self::M as Monoid>::S) -> <Self::M as Monoid>::S {
        return if *f {
            Lengths {
                longest_1_len: x.longest_0_len,
                left_1_len: x.left_0_len,
                right_1_len: x.right_0_len,
                longest_0_len: x.longest_1_len,
                left_0_len: x.left_1_len,
                right_0_len: x.right_1_len,
                len: x.len
            }
        } else {
            x.clone()
        };
    }
    fn composition(f: &Self::F, g: &Self::F) -> Self::F {
        return f^g;
    }
}

G問題

atcoder.jp
今の私の実力で本番これを考察するのはまあ無理ですね。まとめることだけします。
kで場合分けして、総和を取ります。
k=1の場合、f(S,a)-f(S,b)=S_1-S_1=0\lt Xより答えは0です。
k=2の場合、f(S,a)-f(S,b)=S_1a+S_2-S_1b-S_2=S_1(a-b)=Xとなります。
ここから、S_1はXの約数で\displaystyle a-b=\frac{X}{S_1}となることがわかります。また、a\gt bもわかります。
ここで、S_i\lt\min\{10,a,b\}はS_i\lt 10\land S_i\lt a\land S_i\lt bと言い換えられるので、Xの約数かつ10未満であるS_1を固定します。
すると、S_1\lt b\lt a\leq Nとなります。\displaystyle a=b+\frac{X}{S_1}より、これはbの不等式として使うことができ、\displaystyle S_1\lt b\leq N-\frac{X}{S_1}となります。
また、S_2\lt\min\{10,a,b\}です。よって、k=2の場合の答えは\displaystyle\sum_{b=S_1+1}^{N-\frac{X}{S_1}}\begin{cases}10&(b\gt 10)\\b&(b\leq 10)\end{cases}となります。S_1を全探索することを考慮しても時間計算量は\Theta(1)となります。
k\geq 3の場合、\displaystyle f(S,a)-f(S,b)=\sum_{i=1}^k S_ia^{k-i}-\sum_{i=1}^k S_ib^{k-i}=\sum_{i=1}^k S_i(a^{k-i}-b^{k-i})=Xです。a^{k-i}-b^{k-i}はa-bの倍数なので、a-bはXの約数となります。
ここで、a-bはXの約数であるという条件やa\leq Nという条件を考慮しない場合のaとbの組の個数を見積もります。a^2-b^2\leq Xという関係に注目します。
b\lt a\leq\sqrt{X+b^2}より、aとbの組の個数は\displaystyle\sum_{b\geq 1}\left(\left\lfloor\sqrt{X+b^2}\right\rfloor-b\right)となります。
これを不等式評価すると、\displaystyle\sum_{b\geq 1}\left(\left\lfloor\sqrt{X+b^2}\right\rfloor-b\right)\leq\sum_{b=1}^X\left(\sqrt{X+b^2}-b\right)=\sum_{b=1}^X\frac{X}{b\left(1+\sqrt{1+\frac{X}{b^2}}\right)}\leq\sum_{b=1}^X\frac{X}{2b}となり、かなり大まかな評価ですがO(X\log{X})がいえます。
ここで、いささか唐突ですが\displaystyle (a^k-b^k)-(b-1)\sum_{i=0}^{k-1}(a^i-b^i)を計算してみます。
\displaystyle\begin{align*}(a^k-b^k)-(b-1)\sum_{i=0}^{k-1}(a^i-b^i)&=(a^k-b^k)-(b-1)\left(\frac{a^k-1}{a-1}-\frac{b^k-1}{b-1}\right)\\&=(a^k-b^k)+(b^k-1)-(b-1)\frac{a^k-1}{a-1}\\&=(a^k-1)-(b-1)\frac{a^k-1}{a-1}\\&=( (a-1)-(b-1) )\frac{a^k-1}{a-1}\\&=(a^k-1)-(b-1)\frac{a^k-1}{a-1}\\&=(a-b)\frac{a^k-1}{a-1}>0\end{align*}
となり、\displaystyle a^k-b^k>(b-1)\sum_{i=0}^{k-1}(a^i-b^i)がいえます。
S_i\leq b-1であることを考慮すると、Sを決めるには先頭の項から数を貪欲に1つずつ選ぶしかないことがこの不等式からいえます。ここで注意すべき点として、k=2の場合と同様にS_kのみ0以上\min\{10,b\}未満の数を自由にとれる点があります。したがって、aとbの組に対して条件を満たすSの個数はそれぞれ0個か\min\{10,b\}個であるということがわかります。
そして、判定も\Theta(\log X)でできることがわかります。
よって、全体でO(X\log^2{X})で答えを求めることができ、a-bがXの約数である必要があることなどからX\leq 2\times 10^5にしては十分高速に求まります。
atcoder.jp

use ac_library::ModInt998244353;
use num_integer::Roots;

fn main() {
    macro_rules! input {
        ($($tt:tt)*) => {
            proconio::input!($($tt)*);
        };
    }

    #[allow(unused_macros)]
    macro_rules! outputln {
        ($var:expr) => {
            println!("{}",$var);
        };
        ($var:expr,$($vars:expr),+) => {
            print!("{} ",$var);
            outputln!($($vars),+);
        };
    }

    input!(n:usize,x:usize);
    let mut ans=ModInt998244353::raw(0);
    for s1 in 1..10 {
        if x%s1>0 {
            continue;
        }
        if n<x/s1+s1+1 {
            continue;
        }
        ans+=s2(n-x/s1)-s2(s1);
    }
    for b in 2..=min((x-1)/2,n) {
        for a in b+1..=min((x+b*b).sqrt(),n) {
            if x%(a-b)>0 {
                continue;
            }
            let mut apk=1;
            let mut bpk=1;
            while x/(apk*a-bpk*b)>0 {
                apk*=a;
                bpk*=b;
            }
            let mut ok=true;
            let mut x=x;
            while x>0 {
                if x/(apk-bpk)>=min(10,b) {
                    ok=false;
                    break;
                }
                x%=apk-bpk;
                apk/=a;
                bpk/=b;
            }
            if ok {
                ans+=min(10,b);
            }
        }
    }
    outputln!(ans);
}

fn s2(bmax: usize) -> usize {
    return if bmax>10 {
        bmax*10-45
    } else {
        bmax*(bmax+1)/2
    }
}

#[allow(dead_code)]
fn min<T>(left: T, right: T) -> T where T: std::cmp::PartialOrd {
    return if left<right {
        left
    } else {
        right
    };
}