AtCoder Beginner Contest 337 (ABC337) 復習

Rustで競プロするのに慣れたりライブラリを整備したりするのも兼ねて、ABC337を解き直します。
コンテスト時はA～Eの5完でした。
たまにある、他人よりできている感覚はないのになぜか温まっている回でした。風邪気味で凡ミスが多かったのも嫌な感覚の原因でしょう。

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

A問題

atcoder.jp
コンテスト中は入力の順番を間違えてタイムロスしました。
自分はクエリ問題ではない限り大体最初に入力を受け取っていますが、こういう問題はループの中でinput!マクロを書いたほうが良いのかもしれません。

ところで、このスクリプト痒い所に手が届きますね。入れました。
B問題みたいな文字列問題でCopyが大量に出てくるのはちょっと気が散る気もしますが、それでもメリットは大きい気がします。

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 xsum=0;
    let mut ysum=0;
    for _ in 0..n {
        input!(x:usize,y:usize);
        xsum+=x;
        ysum+=y;
    }
    if xsum>ysum {
        outputln!("Takahashi");
    } else if xsum<ysum {
        outputln!("Aoki");
    } else {
        outputln!("Draw");
    }
}

B問題

atcoder.jp
この問題も最初A, B, Cを全て含まないといけないと誤読していました。そのせいでランレングス圧縮をしてしまい、その方針から転換できず鬼の場合分けに。
制約をよく見るとA, B, Cしか登場しないので、ソートするだけで十分でした。
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!(s:Chars);
    let mut ssort=s.clone();
    ssort.sort();
    output_yes_or_no(s==ssort);
}

どうして業プロなら正規表現を真っ先に思いつくはずなのに競プロだと避けてしまうのか…競プロでないときにRustでRegexクレートを触ったことはあったので、使い方はわかっているのに…
そして、最初の誤読を直すのにも+を*に直すだけで済んだので、やっぱりこの方針にすべきでしたね…
atcoder.jp

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

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

    input!(s:String);
    let re=Regex::new(r"^A*B*C*$").unwrap();
    output_yes_or_no(re.is_match(&s));
}

C問題

atcoder.jp
こういう問題は他の問題の小問題として解くことがあるので苦労しなかったです。オアシスですね。
atcoder.jp

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

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

    input!(n:usize,a:[Isize1;n]);
    let mut back=vec![n;n];
    let mut first=n;
    for i in 0..n {
        if a[i]>=0 {
            back[a[i] as usize]=i;
        } else {
            first=i;
        }
    }
    let mut cur=first;
    let mut ans=vec![first+1];
    while back[cur]<n {
        cur=back[cur];
        ans.push(cur+1);
    }
    ans.outputln();
}

D問題

atcoder.jp
累積和を使うべきだというのは察していたのに、それを使わない方針で書いてしまったのはタイムロスだと思います。
道理でみんな4完が早かったわけだ。
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!(h:usize,w:usize,k:usize,s:[Chars;h]);
    let mut ans=usize::MAX;
    let mut dot_sum=vec![0;w+1];
    let mut o_dot_sum=vec![0;w+1];
    for i in 0..h {
        for j in 0..w {
            dot_sum[j+1]=dot_sum[j]+(s[i][j]=='.') as usize;
            o_dot_sum[j+1]=o_dot_sum[j]+(s[i][j]!='x') as usize;
        }
        for j in k..=w {
            if o_dot_sum[j]-o_dot_sum[j-k]==k {
                ans.chmin(dot_sum[j]-dot_sum[j-k]);
            }
        }
    }
    let mut dot_sum=vec![0;h+1];
    let mut o_dot_sum=vec![0;h+1];
    for i in 0..w {
        for j in 0..h {
            dot_sum[j+1]=dot_sum[j]+(s[j][i]=='.') as usize;
            o_dot_sum[j+1]=o_dot_sum[j]+(s[j][i]!='x') as usize;
        }
        for j in k..=h {
            if o_dot_sum[j]-o_dot_sum[j-k]==k {
                ans.chmin(dot_sum[j]-dot_sum[j-k]);
            }
        }
    }
    ans.output_val_or(usize::MAX);
}

E問題

atcoder.jp
Rustに移行してライブラリを整備してローカルの展開のプログラムも追加して、という環境でおそらく初めてのインタラクティブだったので、焦りました。
そして、自分のライブラリのbit_digits関数の仕様で混乱。
他にも、ローカルの展開のプログラムの関係でソースコードだけコピーしてCE、入出力を例まで見ておらず勘違いしていてTLE、など反省が多かったです。

公式解説通りの考察をすぐにできたのだけは救いですかね。
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_interactive!($($tt)*);
        };
    }

    input!(n:usize);
    let m=(n-1).bit_digits()+1;
    outputln!(m);
    for i in 0..m {
        let mut ans=Vec::new();
        for j in 0..n {
            if (j>>i)%2>0 {
                ans.push(j+1);
            }
        }
        print!("{} ",ans.len());
        ans.outputln();
    }
    input!(hint:Chars);
    let mut ans=0;
    for i in 0..m {
        ans+=hint[i].from_fig()*(1<<i);
    }
    ans+=1;
    outputln!(ans);
}

F問題

atcoder.jp
直感的であるだけでなく直感的に解けそうな問題だとは思いつつ、Gのほうが解かれていて自分でも解けそうということから、こちらは考えずG問題にいってしまいました。
蓋を開けてみれば、実際こちらのほうが直感的に解ける簡単な問題でした。

まずCを繰り返すことで、長さ2Nの列の長さNの区間をとって操作を繰り返す問題に落とし込めます。
この前提のもと、それぞれの箱に最初に入るボールについて考えます。
これは、区間に含まれるある色のボールのなかで最初から1\ (\!\!\!\!\!\mod K)個目のもので、かつそのような候補のうち最初からM番目までのものです。
ここで、区間[l,r]をとり、区間に含まれる色cのボールの個数をf(c)とすると、区間に含まれるそのような候補の個数は\sum_c \lceil f(c)/K\rceilであることがわかります。
また、lに対して個数がM以上になる最小のrを考えると、明らかにl_1\lt l_2\Rightarrow r_1\leq r_2が成り立つので、rは尺取り法で求まります。

また、lに対するrがわかると問題の答えもそのままわかります。
各色について箱に最初に入るボールの個数がわかれば、その色のボールは多くともそのK倍の個数しか入らず、またその色のボールの個数より多く入ることもないからです。
よって、答えはf(c)およびCに含まれる色cのボールの個数g(c)を用いて、\sum_c \min\{K\lceil f(c)/K\rceil,\ g(c)\}となります。
gは前計算し、尺取り法を行いながらfおよび2つの式の値を変更していくことで、時間計算量のオーダーは\Theta(N)となります。

実装について、天井関数はdiv_ceil関数を使用しました。また、repeat関数を何気に初めて知りました。
尺取り法のライブラリは作ってありますが、今回は尺取りをしながら多くの変数の値を変更していく必要があるため、自力で実装すべきでしょう。

atcoder.jp

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

#[allow(unstable_name_collisions)]

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

    input!(n:usize,m:usize,k:usize,c:[Usize1;n]);
    let c=c.repeat(2);
    let mut cnts=vec![0;n];
    for i in 0..n {
        cnts[c[i]]+=1;
    }
    let mut idx=0;
    let mut pcnts=vec![0;n];
    let mut boxes=0;
    let mut ans=0;
    for i in 0..n {
        while idx<n+i && boxes<m {
            boxes-=pcnts[c[idx]].div_ceil(&k);
            ans-=min(pcnts[c[idx]].div_ceil(&k)*k, cnts[c[idx]]);
            pcnts[c[idx]]+=1;
            boxes+=pcnts[c[idx]].div_ceil(&k);
            ans+=min(pcnts[c[idx]].div_ceil(&k)*k, cnts[c[idx]]);
            idx+=1;
        }
        outputln!(ans);
        boxes-=pcnts[c[i]].div_ceil(&k);
        ans-=min(pcnts[c[i]].div_ceil(&k)*k, cnts[c[i]]);
        pcnts[c[i]]-=1;
        boxes+=pcnts[c[i]].div_ceil(&k);
        ans+=min(pcnts[c[i]].div_ceil(&k)*k, cnts[c[i]]);
    }
}

G問題

atcoder.jp
既出やんけ！
順位表メタ読みの要素の1つに最速正解時刻があることを知りました。今後使うことになるかはわからないけれども。

コンテスト中は全方位木DPかオイラーツアーかを疑っていました。（早くオイラーツアーをライブラリに追加しないとなぁ…）
TTPC2019の解説では全方位木DPに近い形だったので、そちらを参考に実装も試してみたのですが、諦めました。
公式解説の方法のほうが面白いうえに楽に実装できました。

uと(v,w)の組が条件を満たすことと、uとvとのパスがwを含むことは同値です。そして、これは木からwとwから出る辺を取り除いた森においてvとuが非連結であることとも言い換えられます。
そして、木の根を適当に固定すると、(v,w)の組を決めたときに条件を満たすuの集合は部分木をなすか元のグラフからある部分木を取り除いた木になるかの2通りしかありません。（根つき木を部分木か部分木以外かに分けて考えるのは地味に典型ですね。）

適当な頂点pを決め、pを根とする部分木とそれを取り除いた木について考えます。部分木の頂点の集合をT、取り除いた木の頂点の集合をT^cとします。
(v,w)の組を決めたときに条件を満たすuの集合がTになる条件は、w=p\land v\notin Tとなります。
したがって、そうなる(v,w)の組の個数はw-1-|\{x\in T\ :\ x\lt p\}|となります。
また、\mathrm{parent}(p)をpの親として、(v,w)の組を決めたときに条件を満たすuの集合がT^cになる条件は、w=\mathrm{parent}(p)\land v\in Tとなります。
したがって、そうなる(v,w)の組の個数は|\{x\in T\ :\ x\lt \mathrm{parent}(p)\}|となります。

よって、部分木Tと整数qについて|\{x\in T\ :\ x\lt q\}|を高速に求めることができればこの問題を解くことができます。
DFSの行きがけ順の列を考えると、部分木の頂点は区間をなすので、2次元矩形カウントクエリで解くことができます。
今回は点追加クエリがないため、平面走査とBITを用い、区間の端どうしの差分をとることで求めることができます。
また、TやT^cについて一斉に個数を足し合わせるためには、木上でimos法を応用すればよいです。

意外と苦労せずに実装でき、面白かったです。
atcoder.jp

use ac_library::FenwickTree;
#[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,uv:[(Usize1,Usize1);n-1]);
    let g=VecGraph::construct_graph(n, n-1, &uv);
    let mut ft=FenwickTree::new(n,0);
    let mut imos=vec![0;n];
    let mut seen=vec![false;n];
    dfs(n, &g, 0, n, &mut ft, &mut imos, &mut seen);
    let mut ans=vec![0;n];
    ans[0]=imos[0];
    for gs in g.dfs(0) {
        match gs {
            GraphSearch::Vertex(_, _) => (),
            GraphSearch::VertexEdgeWeight(v, u, _) => ans[u]=ans[v]+imos[u]
        }
    }
    ans.outputln();
}

fn dfs(n: usize, g: &VecGraph, w: usize, par: usize, ft: &mut FenwickTree<usize>, imos: &mut Vec<isize>, seen: &mut Vec<bool>) {
    seen[w]=true;
    let prev_w=ft.sum(0..w);
    let prev_par=ft.sum(0..par);
    ft.add(w, 1);
    for &(v,_) in &g.get()[w] {
        if !seen[v] {
            dfs(n, g, v, w, ft, imos, seen);
        }
    }
    let cnt_w=w-ft.sum(0..w)+prev_w;
    let cnt_par=ft.sum(0..par)-prev_par;
    imos[w]+=cnt_w as isize;
    imos[0]+=cnt_par as isize;
    imos[w]-=cnt_par as isize;
}