AtCoder Beginner Contest 335 (ABC335) 復習

Rustで競プロするのに慣れたりライブラリを整備したりするのも兼ねて、ABC335を解き直します。
コンテスト時はA~Dの4完でした。
いよいよもって落水の危機に立たされました。

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

A問題

atcoder.jp
こういう文字列操作の問題は未だにどうRustで書くべきか迷いますね。
スライスを使うときは&strにしなければならないことに注意。
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!(s:String);
    println!("{}4",&s[0..s.len()-1]);
}

B問題

atcoder.jp
要求されていることも問題文もめちゃくちゃ単純なのにめちゃくちゃ誤読しました。
最近入出力例を最初に読まない弊害がここにも。
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);
    for i in 0..=n {
        for j in 0..=n {
            for k in 0..=n {
                if i+j+k<=n {
                    outputln!(i,j,k);
                }
            }
        }
    }
}

C問題

atcoder.jp
VecDequeではなくVecでもできるとは本番中も予想していましたが、頭を使いたくなかったのでVecDequeでやりました。
(isize,isize)で用意すべきか2つのVecDequeを用意すべきか、ifで書くべきかmatchで書くべきかなどを検討してみましたが、個人的には本番中と同じ書き方がよいと感じたのでそのままで突き通します。
atcoder.jp

use std::collections::VecDeque;

#[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,q:usize);
    let mut d=VecDeque::<(isize,isize)>::new();
    for i in 0..n {
        d.push_back((i as isize+1,0));
    }
    for _ in 0..q {
        input!(q:usize);
        if q==1 {
            input!(c:char);
            if c=='R' {
                d.push_front((d[0].0+1,d[0].1));
            } else if c=='L' {
                d.push_front((d[0].0-1,d[0].1));
            } else if c=='U' {
                d.push_front((d[0].0,d[0].1+1));
            } else {
                d.push_front((d[0].0,d[0].1-1));
            }
            d.pop_back();
        } else {
            input!(p:Usize1);
            outputln!(d[p].0,d[p].1);
        }
    }
}

D問題

atcoder.jp
本番中はVecで文字列を持つようにしておけば出力が楽になるということに気付いていませんでした。
そして、間違えて全てprintln!で出力して、提出されてしまってからそれに気付き焦りました。ACになる仕様でよかった。
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);
    let mut d=vec![vec![String::new();n];n];
    let mut i=0;
    let mut j=0;
    let mut dir=0;
    let mut cnt=1;
    while i!=n/2 || j!=n/2 {
        d[i][j]=format!("{}",cnt);
        cnt+=1;
        if dir==0 {
            j+=1;
            if j==n-1 || d[i][j+1].len()>0 {
                dir=1;
            }
        } else if dir==1 {
            i+=1;
            if i==n-1 || d[i+1][j].len()>0 {
                dir=2;
            }
        } else if dir==2 {
            j-=1;
            if j==0 || d[i][j-1].len()>0 {
                dir=3;
            }
        } else {
            i-=1;
            if i==0 || d[i-1][j].len()>0 {
                dir=0;
            }
        }
    }
    d[n/2][n/2]="T".to_string();
    d.outputlns();
}

E問題

atcoder.jp
苦しかった。ですが、最長経路問題、強連結成分分解、ダイクストラ法などについての理解が足りなかったのは事実だと思います。
数値の等しい隣接頂点の同一視やDAG上のDPによる最長経路問題の求解などは当然浮かんだのですが、入る頂点の数値が出る頂点の数値以上であるような有向グラフを考えればサイクルができるのは数値が等しい隣接頂点のみであることに気付けず、またACLのSCCがトポロジカルソートまでしてくれることを知りませんでした。
それらがわかっていれば、強連結成分を頂点とするDAG上でDPをするだけでした。

本番中は結局嘘ダイクストラ法に嵌り、そのままTLEが取れずに終了しました。
精進はまだまだ必要ですね。
atcoder.jp

use ac_library::SccGraph;
#[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,a:[usize;n],uv:[(Usize1,Usize1);m]);
    let g=VecGraph::construct_graph(n, m, &uv);
    let mut sg=SccGraph::new(n);
    for &(u,v) in &uv {
        if a[u]<=a[v] {
            sg.add_edge(u, v);
        };
        let (u,v)=(v,u);
        if a[u]<=a[v] {
            sg.add_edge(u, v);
        };
    }
    let scc=sg.scc();
    let mut ind=vec![n;n];
    for i in 0..scc.len() {
        for &v in &scc[i] {
            ind[v]=i;
        }
    }
    let mut dp=vec![0;scc.len()];
    let mut seen=vec![false;scc.len()];
    dp[ind[0]]=1;
    seen[ind[0]]=true;
    for i in ind[0]..scc.len() {
        if seen[i] {
            let dist=dp[i];
            for &v in &scc[i] {
                for &(u,_) in &g.get()[v] {
                    if ind[u]>i {
                        seen[ind[u]]=true;
                        dp[ind[u]].chmax(dist+1);
                    }
                }
            }
        }
    }
    outputif(seen[ind[n-1]], dp[ind[n-1]], 0);
}

F問題

atcoder.jp
本番中は公式解説の素朴なDP(最も右の黒いマスがiであるときの場合の数を求める、前からのDP)を思いつかず、後ろから総和をとらない形で考えていたので、そのまま考えていてもそもそも時間内に計算量の削減までたどり着けたかわかりません。まあそれが実力ではあると思います。

\mathrm{DP}[i]を「最も右の黒いマスがiであるときの、黒く塗られたマスの集合の個数」とすると、次にたどり着けるマスの値に今のマスの値を足す素朴なDPを考えることができます。しかし、これではA_iが小さいときに間に合いません。
そこでA_iで割った余りに注目します。(これは一応本番中の私も考えていました。)
DPの漸化式は次にたどり着ける全てのマス、すなわちA_iの倍数だけ進んだ全てのマスに今のマスの値を足すものです。ゆえに、\mathrm{DP}_2[a][m]を「A_i=ai\equiv m\ (\!\!\!\!\!\mod A_i)であるDPテーブルの総和」としておき、次にたどり着けるマスにたどり着いた場合に丸ごと足せば計算量を抑えられます。
そして、2つの方法の境界を適切にとることで時間計算量のオーダーはO(N\sqrt{N})となります。

atcoder.jp

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

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

    input!(n:usize,a:[usize;n]);
    let mut dp=vec![Mint::new(0);n];
    dp[0]+=1;
    let s=n.sqrt()+1;
    let mut dp2=vec![vec![Mint::new(0);s];s];
    for i in 0..n {
        for m in 1..s {
            dp[i]+=dp2[m][i%m];
        }
        if a[i]>=s {
            for j in (i+a[i]..n).step_by(a[i]) {
                dp[j]=dp[j]+dp[i];
            }
        } else {
            dp2[a[i]][i%a[i]]+=dp[i];
        }
    }
    outputln!(dp.sum());
}

Kiri8128さんのユーザ解説のほうも挑戦しました。
実装が簡単なうえに、計算量解析が非自明で平方分割と同じオーダーになるのがとても面白いと思います。

素朴なDPの枝刈りとして、A_i=A_jかつ進めるマスがずれていない(i,j)の組があれば、ijより後の更新をjのときにまとめてやってしまうというものが考えられます。
実はこれをするだけで時間計算量のオーダーはO(N\sqrt{N})に落ちます。
G_{k,m}=\{i\ |\ i\equiv m\ (\!\!\!\!\!\mod A_i)\},\ f(G_{k,m})=N/kとおくと、ループ数は\sum_{G_{k,m}\neq\varnothing}f(G_{k,m})で抑えられます。
これを最大化するAは、1,2,2,3,3,3,4,4,4,4,\ldotsという形になります。
そして、N三角数のとき、同じ数でのfの和がNになることから、ループ数はN\max{A}=O(N\sqrt{N})で抑えられます。
また、無理関数が上に凸である性質からN三角数でなくてもループ数はO(N\sqrt{N})で抑えられます。

実装について、自分は\mathrm{DP}_2[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 dp=vec![Mint::new(0);n];
    let mut dp2=vec![Mint::new(0);n];
    dp[0]+=1;
    dp2[0]+=1;
    for i in 0..n-1 {
        for j in (i+a[i]..n).step_by(a[i]) {
            dp[j]=dp[j]+dp2[i];
            dp2[j]=dp2[j]+dp2[i];
            if a[i]==a[j] {
                dp2[j]=dp2[j]+dp2[i];
                break;
            }
        }
    }
    outputln!(dp.sum());
}

G問題

橙diffのG問題は自分にはまだ早いというのはもちろん当然として、解けている人数のわりに群論をちゃんと勉強していれば難しくない内容だったのでいつか解けるようになりたいです。

{A_i}^k\equiv A_j\!\!\!\!\mod Pとなる(i,j)の個数を求める問題です。
P素数なので、乗法群(\mathbb{Z}/P\mathbb{Z})^*巡回群C_{P-1}と同型です。
有限巡回群について、以下のような性質があります。

  • 1つの元が生成する有限巡回群は元の有限巡回群の部分群となり、その元の位数は生成する群の位数である。
  • 有限巡回群の任意の部分群は有限巡回群であり、部分群の位数は元の群の位数の約数となる。
  • 有限巡回群の位数の任意の約数について、位数がその約数となる部分群がただ1つ存在する。

これらの性質から、{A_i}^k\equiv A_j\!\!\!\!\mod P\langle A_j\rangle\subseteq\langle A_i\rangle\subseteq(\mathbb{Z}/P\mathbb{Z})^*と同値であり、さらにこれは|\langle A_j\rangle|\Big{|}|\langle A_i\rangle|と同値であることがわかります。

したがって、まず全ての元の位数を求めればよいです。元の位数は、例えば素因数分解をしたうえで試し割りと繰り返し二乗法を用いることでO({\log}^2 P)で求められます。
また、位数として考えられる個数は少ないので、位数の重複をまとめたうえで組をとり、片方が片方を割り切るかを愚直に確認することで解くことができます。

PAの上限が大きく、繰り返し二乗法で32bitの範囲に収まらない数の2乗を計算する必要があります。
BigUintのmodpow関数は重くTLEしたので、u128を使うmodpow関数をライブラリに追加しました。
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,p:usize,a:[usize;n]);
    let pfs=(p-1).prime_factorize_using_pollard_s_rho();
    let mut order=vec![p-1;n];
    for i in 0..n {
        for &(pf,_) in &pfs {
            while order[i]%pf==0 {
                if a[i].modpow(order[i]/pf, p)==1 {
                    order[i]/=pf;
                } else {
                    break;
                }
            }
        }
    }
    order.sort();
    let os=order.rle();
    let l=os.len();
    let mut ans=0;
    for i in 0..l {
        for j in 0..l {
            if os[i].0%os[j].0==0 {
                ans+=os[i].1*os[j].1;
            }
        }
    }
    outputln!(ans);
}