AtCoder Regular Contest 169 (AからCまで) 復習

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

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

A問題

atcoder.jp
難しかったですが面白かったです。
コンテスト中は1回雑嘘解法でWAを出してから、B問題をACしたことにより余裕が出て正解にたどり着けた気がします。
とはいえ、簡単な例で操作後の値を計算してみるのを面倒臭がって途中までやらなかったのは良くなかった気がします。

明らかに最終的なA_1は最初の時点でのA_iの線形結合になります。そして、P_i\lt iの制約から、実験するとi\to P_iのグラフを考えたときの頂点1までの距離が大きいほうが係数が大きいことがわかります。そして、距離が等しければ係数は等しいので、距離が大きい順に等しい距離のA_iの和をとり、0ならば保留して、そうでなければその正負を答えとすればよいです。距離が0の場合まで和が全て0であるときに限り答えは0です。(前回のARCと今回のARCはWriterが同じmaroonrkさんで、前回のB問題も答えが決まらないならば保留しつつ大きいほうから順に見るという構造だったので思いつけたのもあります)

コンテスト中はこれだけ考えて通しましたが、厳密に係数を評価します。
まず、操作前のA_iのベクトルを\boldsymbol{A}、操作後のA_iのベクトルを\boldsymbol{A}'とします。すると、M_{i,i}M_{P_i,i}のみが1で残りは0の行列Mを用いて、\boldsymbol{A}'=M^{10^{100}}\boldsymbol{A}と表せます。ここで、{M'}_{P_i,i}のみが1で残りは0の行列M'を考えると、\boldsymbol{A}'=(I+M')^{10^{100}}\boldsymbol{A}と変形できます。(ただしI単位行列
単位行列との行列積は可換ですから、\boldsymbol{A}'=\sum_{k=0}^{10^{100}}\binom{10^{100}}{k}{M'}^k\boldsymbol{A}というふうに二項定理で展開することができます。
ここで、{P_i}^1=P_i,\ {P_i}^{k+1}=P_{{P_i}^k}とおいておきます。たとえば(M')^kに全ての成分が1のベクトルをかけたときを考えると、(M')^k{P_i}^k,iの成分のみが1で残りは0の行列となります。したがって、i\to P_iのグラフを考えたときの頂点1までの距離がkならば係数は\binom{10^{100}}{k}となります。\binom{10^{100}}{k}/\binom{10^{100}}{k-1}A_iよりも明らかに大きいので、距離の大きさだけを考慮すればよいことがわかります。

実装について、コンテスト中はBTreeMapを使ってしまいましたが、普通にVecを使ったほうが実装量も速度も良いですね。精進。
atcoder.jp

#[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,a:[isize;n]);
    let mut p=vec![1;n];
    for i in 1..n {
        input!(tmp:Usize1);
        p[i]=tmp;
    }
    let mut dist=vec![0;n];
    let mut dsum=vec![0;n];
    dist[0]=0;
    dsum[0]+=a[0];
    for i in 1..n {
        dist[i]=dist[p[i]]+1;
        dsum[dist[i]]+=a[i];
    }
    for &s in dsum.iter().rev() {
        if s>0 {
            outputln!("+");
            return;
        } else if s<0 {
            outputln!("-");
            return;
        }
    }
    outputln!(0);
}

B問題

atcoder.jp
区間に対して定まる値の全区間についての総和を求める問題、得意です。
コンテスト中にfをまず1つ求めるときに貪欲でいいかを5分くらい悩んでしまいました。
貪欲法の証明は、貪欲をしない方法から貪欲をする方法に変えたときに損をしないことをいうのが定石らしいですね。
逆方向で言おうとしていたのもあって、腑に落ちなかったのかもしれません。

解説とは少し異なるDPで解いたかもしれません。
lについて、f([l,r))=1となる最大のrは尺取り法で求まります。
そしてf([l,i])の値は、iが貪欲に区切った同じ区間に含まれていれば等しく、右の区間ほど大きくなっていきます。
また、貪欲に区切った2番目以降の区間は、その左端から貪欲に区切った区間と全く同じになります。
そこで、同じ区間に含まれているものを積で計算し、途中から同じ区切り方になるところも積で計算することを考えます。
簡単に区間を2つに区切ることができる場合を考えると、最初の区間は1×区間の幅、次の区間はそれを1番目とするものも含めて3×区間の幅となります。
一般化すると、c個の区切り方で左端になりその寄与がwである場合、この区間の次の区間(c+1)個の区切り方で左端になりその寄与はw+cになります。
これを尺取り法とDPで実装すればよいです。

実装では尺取り法の境界でちょっとだけ添字ガチャをしてしまいました。翌日usize用の二分探索の関数で事故ったのもあり今は仕様を変えています。
添字ガチャ中にループ内で全てのベクター標準エラー出力を書いたまま提出してしまいましたが、自作の関数とマクロはデバッグビルドでしか出力しないようにしていたので事なきを得ました。
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,a:[usize;n]);
    let asum=a.construct_prefix_sum();
    let mut w=vec![1;n+1];
    let mut c=vec![1;n+1];
    let mut ans=0;
    for (l,r) in two_pointers(n, n+1, true, true, &|l,r| asum.calculate_partial_sum(l, r)<=s) {
        ans+=w[l]*(r-l);
        w[r]+=w[l]+c[l];
        c[r]+=c[l];
    }
    outputln!(ans);
}

C問題

atcoder.jp
コンテスト中、公式解説とは異なるqueueを使う方法で方針を思いついていましたが、A問題に手こずっていたのもあり実装が間に合いませんでした。
とはいえ、黄diffに対してコンテスト中の方針が合っていたのは嬉しいです。

\mathrm{DP}[i][j][k]を、i番目の要素で最後にjk個連続しているときの場合の数というDPはすぐ思いつきます。
空間計算量は\mathrm{DP}[j][k]とすればよいだけなので、時間計算量のオーダーを小さくすることを考えましょう。
思い切って、このままで漸化式を考えてみます。
A_i=-1のとき、全てのjについて(j+1)個にならない限りそのまま繋げられるので、DPテーブルはキューのように全てずれていきます。
k=0は、j以外の全ての場合の数をそのまま乗せればよいです。
A_i\neq -1のとき、A_i=jならばA_i=-1と同じように考えてよいです。
A_i\neq jならば、全てのテーブルは0になります。
これらを踏まえると、\sum_{j=1}^N\sum_{k=1}^j\mathrm{DP}[j][k],\ \sum_{k=1}^j\mathrm{DP}[j][k]と、全てのテーブルを0にするための「今DPテーブルで正しい値を持っているのは何個目までか」という値を別に持てばいいことがわかります。

これらを実装すればよいです。最初の1通りに注意。
atcoder.jp

use std::collections::VecDeque;

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,a:[isize;n]);
    let mut dp=vec_range(0, n+1, |i| VecDeque::<ModInt998244353>::from(vec![ModInt998244353::new(0);i]));
    let mut sum=ModInt998244353::new(1);
    let mut sums=vec![ModInt998244353::new(0);n+1];
    let mut valid=vec_range(0, n+1, |i| i);
    for i in 0..n {
        let prev_sum=sum;
        if i==0 {
            sum-=1;
        }
        if a[i]==-1 {
            for j in 1..=n {
                dp[j].push_front(prev_sum-sums[j]);
                let tmp=sums[j];
                sums[j]=prev_sum;
                sum+=prev_sum-tmp;
                if valid[j]==j {
                    sums[j]-=dp[j].back().unwrap();
                    sum-=dp[j].back().unwrap();
                } else {
                    valid[j]+=1;
                }
                dp[j].pop_back();
            }
        } else {
            for j in 1..=n {
                if a[i] as usize==j {
                    dp[j].push_front(prev_sum-sums[j]);
                    let tmp=sums[j];
                    sums[j]=prev_sum;
                    sum+=prev_sum-tmp;
                    if valid[j]==j {
                        sums[j]-=dp[j].back().unwrap();
                        sum-=dp[j].back().unwrap();
                    } else {
                        valid[j]+=1;
                    }
                    dp[j].pop_back();
                } else {
                    sum-=sums[j];
                    sums[j]=ModInt998244353::new(0);
                    valid[j]=0;
                }
            }
        }
    }
    outputln!(sum);
}

D以降は流石に私には厳しそうです。絶妙に他のやりたいこともありますし。