【Rust】ABC417まとめ（A~E）

AtCoder Beginner Contest 417

今回は15分遅れでunrated参加しました。
A~Cまでの3完で、コンテスト後にD,Eをupsolveしたので、
今回はA~Eまでまとめます。

A - A Substring

左から$A$文字、右から$B$文字削除することは、裏を返せば$A+1$文字目から$N-B$文字目までを残すと言い換えることができます。
Rustではインデックスが0から始まるので、$A$文字目から$N-B-1$文字目までです。

文字列$S$の$a$文字目から$n-b-1$文字目までのスライスを取得するためには、s[a..(n-b)]もしくはs[a..=(n-b-1)]のように書きます。
なお、println!マクロに&strを渡すと、内部で自動参照外しが行われstr型となり、str型がサイズ不定型（Dynamically Sized Type）のためコンパイルエラーが発生します。Rustではそのため、あえて&s[a..=(n-b-1)]として&&str型を渡すことでコンパイルエラーを回避する必要があります。

補足

Rustにはderef coercionやauto dereferencing ruleなどのルールがあり複雑に絡み合っているようで、筆者も全てを理解したわけではありません。
gpt-5 thinkingに聞いても頓珍漢な回答が返ってくることが多々あるのですが、その中でも一番的を射てそうな回答を補足として貼っておきます。（正しい確証はないので注意願います）

1. str はサイズ不定型（DST）で 値としては扱えない（Sized を満たさない）。
2. String や &str に対する s[a..b] の 「出力の型の名前」 は str（Index::Output = str）と定義されています。
3. インデックス演算の実装自体は &self を取って &Self::Output を返すので、普通は s[a..b] は &str のように見えます。
4. しかし、コンパイラがその式を 「サイズが必要な文脈」（値として移動したり、Sizedが要求される場所）に置くと str が値扱いになろうとして Sized エラーが発生します。
5. println!("{}", s[a..b]); がエラーになる典型的な理由はまさにここで、コンパイラが str を値として扱う（＝サイズ不定なので拒否）文脈になっているためです。
6. 一方 &s[a..b] を書くと明示的に &str（参照＝ポインタでサイズ確定） を作るので、Sized 要求を満たして安全に渡せます。
7. つまり問題は「str（DST）を値として渡そうとしているかどうか」で、& を付けることで str を参照にして回避する、という話です。
8. 補足：&&str が通るのは、&&str → 自動的に参照を外して &str として Display 等の実装が使われるからです（deref coercion）。
9. 実務的な対処法は単純で、スライスは一旦変数に束縛して let slice: &str = &s[a..b]; println!("{}", slice); のように明示することです。
10. 要点：DST (str) は値にならない → 参照 &str にして渡す → それが安全でコンパイルが通る。

人間が書いた記事としては以下が参考になると思います。
type coercion(型強制)に慣れ親しむ

use proconio::input;

fn main() {
    input! {
        n: usize,
        a: usize,
        b: usize,
        s: String,
    }

    println!("{}", &s[a..(n-b)]);
}

B - Search and Delete

配列$A,B$の要素の範囲は$0 \le A_i, B_i \le 10^9$とかなり広いのですが、$N \le 100$と要素の種類数が小さいのでstd::collections::HashMapを使って個別の値について$\text{cnt}=Aで出現した回数-Bで出現した回数$を記録して、$\text{cnt}$が$1$以上のもののみ配列$\text{arr}$にpushして、昇順にソートした結果を出力すれば良いです。

例えば、問題のサンプル1の場合を見てください。

例）サンプル1
8 5
1 2 2 3 3 3 5 6
2 2 7 3 2

$A$での出現回数と$B$での出現回数、およびその差を表にまとめます。

キー	$A$での出現回数	$B$での出現回数	$A$での出現回数$-B$での出現回数
1	1	0	1
2	2	3	-1
3	3	1	2
5	1	0	1
6	1	0	1
7	0	1	-1

上の結果より、「$A$での出現回数$-$$B$での出現回数」が$1$以上のキーだけ抽出して配列を復元すると$[1,3,3,5,6]$です。これは出力例1と一致します。
なお、std::collections::HashMapを利用する場合は必ずしもキーの小さい順に取得できるとは限らないため、結果配列をソートしておく必要があることに注意してください。

use proconio::input;
use itertools::Itertools;

fn main() {
    input! {
        n: usize,
        m: usize,
        a: [usize; n],
        b: [usize; m],
    }
    let mut map = std::collections::HashMap::new();
    for i in 0..n {
        *map.entry(a[i]).or_insert(0) += 1;
    }

    for i in 0..m {
        if let Some(cnt) = map.get_mut(&b[i]) {
            *cnt -= 1;
            if *cnt == 0 {
                map.remove(&b[i]);
            }
        }
    }

    let mut result = Vec::new();
    for (&key, &count) in &map {
        for _ in 0..count {
            result.push(key);
        }
    }
    result.sort();
    println!("{}", result.iter().join(" "));
}

C - Distance Indicators

$j-i=A_i+A_j$となる整数の二つ組$i,j \ (1\le i,j \le N,i \ne j)$ですが、これを愚直に全探索しようとすると$O(N^2)$となりTLE必至のため、少なくとも$O(NlogN)$以下に効率化しましょうと言う問題です。

まず、$j-i=A_i+A_j$のままではどうしても$i,j$の二重ループで全探索する以外に方法はありませんので、どうにか式変形することでこれを回避できないかと言うことを考えます。

そこで、$j-i=A_i+A_j$を$j-A_j=A_i+i$のように式変形することを考えます。
このようにすることで左辺と右辺どちらもインデックスが揃い、$\text{pl}(i)=i+A_i$、$\text{mi}(i)=i-A_i$のように$i$を引数とする二つの関数が$O(N)$で定義できます。

つまり、問題文は次のように言い換えることができます。
「$\text{pl}(i)=\text{mi}(j)$を満たす整数の二つ組$i,j \ (1\le i,j \le N,i \ne j)$の組み合わせ数を求めよ。ただし、$\text{pl}(i)=i+A_i$、$\text{mi}(i)=i-A_i$とする」と。

この発想ができれば、あとはかなり簡単です。
std::collections::HashMapにて、$\text{pl}(i)(1\le i \le N)$で出現する値と出現回数、$\text{mi}(i)(1\le i \le N)$で出現する値と出現回数をそれぞれ$O(N)$で記録し、キーごとに「$\text{pl}での出現回数 \times \text{mi}での出現回数$」を足し上げていくだけで良いからです。

なお、もし$A_i=0$がありうる場合は$\text{pl}(i)=\text{mi}(i)$も起こり得るのですが、今回は$A_i > 0$のため必ず$\text{pl}(i) \ne \text{mi}(i)$が保証されているのでダブルカウントの問題はありません。

以上、全体として$O(N)$でこの問題を解くことができました。

use proconio::input;

fn main() {
    input! {
        n: usize,
        a: [i64; n],
    }

    let mut pl = vec![0; n];
    let mut mi = vec![0; n];
    for i in 0..n {
        pl[i] = i as i64 + a[i];
        mi[i] = i as i64 - a[i];
    }

    let mut fq_pl = std::collections::HashMap::new();
    let mut fq_mi = std::collections::HashMap::new();
    for i in 0..n {
        *fq_pl.entry(pl[i]).or_insert(0) += 1;
        *fq_mi.entry(mi[i]).or_insert(0) += 1;
    }

    let mut ans: i64 = 0;
    for (key, &pl_count) in &fq_pl {
        if let Some(&mi_count) = fq_mi.get(key) {
            ans += pl_count * mi_count;
        }
    }
    println!("{}", ans);
}

D - Takahashi's Expectation

この問題は、$1 \le P_i, A_i ,B_i \le 500$という制約条件に着目し、二分探索・DP（動的計画法）・累積和という典型的なアルゴリズムを三つ組み合わせて解く必要があるため、いつものD問題に比べて若干難しいと思いました。

まず、$1 \le P_i, A_i ,B_i \le 500$という条件によってどういうことが起こりうるかを考えましょう。
最も大事なのは$1 \le P_i \le 500$です。現在のテンション$X$が$P_i$以上の時テンションは$A_i$上がり、$P_i$より小さい時$B_i$下がると言うふうに、$P_i$を境に高橋くんのテンション遷移が大きく変わるからです。

▼テンション$X$とプレゼントの価値$P_i$の関係別でのテンション遷移イメージ

この条件によって、$i$番目のプレゼントにおいて、テンション$X$が$500$以上ならばおのずと$X \ge P_i$であることが保証され、テンションが$B_i$だけ下がると言うことが確定します。 したがって、テンションが非常に大きい場合はテンションが$500$以下になるまで$B_i$を繰り返し引けば良いのですが、これは二分探索と累積和を駆使して$X-\sum_{i=1}^k B_i \le 500$となるような最初のインデックス$k$を求めることで、テンションが$501$を初めて下回るタイミングと、そのタイミングにおけるテンションの値を$O(logN)$で高速に求めることができます。

一方で、$k$番目のプレゼントでテンションが$501$を下回った後の挙動についてはテンションの大きさとプレゼントの価値の大きさを比較した結果によって大きく変わるため一筋縄ではいきません。
また、$X$が$501$を下回った段階からシミュレーションをするというのも一案かもしれませんが、全ての質問$X_i$が$1000$程度とそれほど大きくない場合、最悪で$NQ$回近くのループを回さなくてはなりませんから、これはTLEとなるリスクが非常に高いです。

そこで、次のように動的計画法を定義します。
$$dp_{i,j}:=現在のテンションがiで、次にj番目のプレゼントをもらうとき、最終的なテンションはいくつになるか$$

現在のテンション$i$と$j$番目のプレゼントをもらった次の瞬間のテンション$ni$については、

• $i \le P_j$のとき：$ni=i+A_j$
• $i > P_j$のとき：$ni=\text{max}(0,i-B_j)$（0を下回ってはいけないため）

ということがすぐにわかります。
ここで、もし$dp_{-,j+1}$($-$は$dp$配列内の任意の値)が全てわかっているとすると、$dp_{ni,j+1}$の中身を見れば最終的なテンションも連鎖的にわかります。したがって、$dp_{i,j}←dp_{ni,j+1}$のように、後ろから前へと$dp$値を伝播していきます。
なお、$ni$の下限は$0$ですが、$i=500,P_i=500,A_i=500$が起こりうる$ni$の最大ケースのため、上限は$1000$です。よって、$dp$テーブルは$1001 \times N$だけ用意しておきましょう。
こうして、初めて$X \le 500$となった瞬間以降の最終的なテンションも$O(N\max (P))$で前計算しておくことができました。

補足

$X \ge 501$のとき必ずテンションは下がるとはいえ、$dp$テーブルのサイズを$501 \times N$などとしてはいけません。あくまで二分探索+累積和でわかるのは初めてテンションが$501$を下回る瞬間であり、その後にテンションが上がって$501$以上になることもありえます。

このとき、もし$dp$の第一引数が$500$以下しか受け付けないようであれば、配列外参照となりRustがpanicを引き起こします。
（補足おわり）

以上で準備が完了し、初期値$X_i$がどのような値であっても高速に最終的なテンションが求められるようになっています。具体的な手順は以下のとおりです。

1. $X_i$が$501$以上の時、$X_i$が初めて$501$を下回るタイミング$i$とその時点でのテンション$j$を求める。もし下回らない場合は$X_i - \sum B$を出力して終了する
2. $dp_{i,j}$を参照し、最終的なテンションを求めて出力する

一つ目の手順は二分探索と累積和を利用するため$O(logN)$です。
二つ目の手順は、あらかじめ$dp$テーブルを前計算しているため$O(1)$で済みます。
したがって、全ての質問のにかかる計算量は$O(QlogN)$であり、累積和と$dpテ$ーブルの前計算にかかる計算量は$O(N\text{max}(P))$であるため、最終的な計算量は$O(N\text{max}(P)+QlogN)$となり、これは十分高速です。

use proconio::input;
use itertools::Itertools;

pub mod binary_search {
    pub fn is_ok(arr: &[usize], key: usize, mid: usize) -> bool {
        key.saturating_sub(arr[mid]) <= 500
    }

    pub fn lower_bound(arr: &[usize], key: usize) -> usize {
        let mut ng = -1isize;
        let mut ok = arr.len() as isize;

        while ok - ng > 1 {
            let mid = (ok + ng) / 2;
            if is_ok(arr, key, mid as usize) {
                ok = mid;
            } else {
                ng = mid;
            }
        }
        ok as usize
    }
}

fn main() {
    input! {
        n: usize,
    }

    let mut p = vec![0usize; n];
    let mut a = vec![0usize; n];
    let mut b = vec![0usize; n];
    for i in 0..n {
        input! {
            pi: usize,
            ai: usize,
            bi: usize,
        }
        p[i] = pi;
        a[i] = ai;
        b[i] = bi;
    }

    let mut dp = vec![vec![0usize; 1001]; n];
    for i in 0..1001 {
        dp[n - 1][i] = if p[n - 1] >= i {i + a[n - 1]} else {i.saturating_sub(b[n - 1])};
    }
    
    let mut sum_b = vec![0usize; n + 1];
    for i in 0..n {
        sum_b[i + 1] = sum_b[i] + b[i];
    }

    for i in (0..n - 1).rev() {
        for j in 0..1001 {
            dp[i][j] = if p[i] >= j {
                dp[i + 1][j + a[i]]
            } else {
                dp[i + 1][j.saturating_sub(b[i])]
            };
        }
    }

    input! {    
        q: usize,
        x: [usize; q],
    }

    let mut ans = vec![];
    for xi in x {
        let idx = binary_search::lower_bound(&sum_b, xi);
        if idx >= n {
            ans.push(xi - sum_b[n]);
        } else {
            ans.push(dp[idx][xi.saturating_sub(sum_b[idx])]);
        }
    }

    println!("{}", ans.iter().join("\n"));
}

E - A Path in A Dictionary

まず、$X$から$Y$までの単純パスは桁違いに多くなる可能性があるため、全列挙して辞書順ソートした後に先頭のパスを出力するという方法はTLEになることを理解する必要があります。

ここで、もっとシンプルに考えると、「$X$を始点として、辞書順に小さいノードから優先してパス記録しながら訪問し、$Y$に到達した瞬間に探索を終了する」というdfs(深さ優先探索)で辞書順最小パスを取得できることに気づくことができます。

これは、文字列の後ろの方よりもなるべく前の方で小さい値を選んだ方が辞書順に小さくなるという性質を考えれば気づけると思います。
あとは、きちんとdfsを実装できるかですが、帰りがけに$\text{visited}$配列をfalseにしてしまうと、以下のようなグラフで辞書順に最小のノードに何度も訪問してしまうという無限ループが発生するので、それはやらないようにしましょう。

▼帰りがけに$\text{visited}$配列をfalseにしてしまうと、$X$から$1$に何度も訪問し無限ループとなる

グラフ$G$は無向グラフのため、一度訪問したノードの先に$Y$に到達できないのであれば、次にそのノードを訪問することは無用のため、$\text{visited}$配列をfalseにして再訪問できるようにしなくてもよいのです。

use proconio::input;
use itertools::Itertools;
fn main() {
    input! {
        t: usize,
    }

    fn dfs(g: &Vec<Vec<usize>>, visited: &mut Vec<bool>, cur: usize, goal: usize, st: &mut Vec<usize>) -> bool {
        visited[cur] = true;
        st.push(cur + 1);
        if cur == goal {
            return true;
        }
        for &next in &g[cur] {
            if visited[next] {
                continue;
            }
            if dfs(g, visited, next, goal, st) {
                return true;
            }
        }
        st.pop();
        false
    }

    for _i in 0..t {
        input! {
            n: usize,
            m: usize,
            x: usize,
            y: usize,
        }

        let mut u = vec![0usize; m];
        let mut v = vec![0usize; m];
        let mut g = vec![vec![]; n];
        for j in 0..m {
            input! {
                uij: usize,
                vij: usize,
            }
            u[j] = uij;
            v[j] = vij;
            g[uij - 1].push(vij - 1);
            g[vij - 1].push(uij - 1);
        }
        for j in 0..n {
            g[j].sort();
        }
        let mut visited = vec![false; n];
        let mut st = vec![];
        dfs(&g, &mut visited, x - 1, y - 1, &mut st);
        println!("{}", st.iter().join(" "));
    }
}

まとめ

Rustを書き始めて3週間くらいですが、だんだんと書き慣れてきました。
F問題は遅延セグメント木のライブラリコピペが必要で、いまだにAtCoder関連のライブラリの使い方がわからないため解くことができませんでした。
今後、ライブラリの使い方にも慣れてきたらRated参加してみようと思います。