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

AtCoder Beginner Contest 418

unratedで参加しました。
A~Dの4完（89分1ペナ）で、コンテスト後にEをupsolveしました。
今回はA~Eまでまとめます。

A - I'm a teapot

文字列の後ろ3文字をスライスして、teaと一致するかどうかを調べます。
Rustではs[a..b]で文字列のa文字目からb-1文字目までの文字列スライスを取得でき、bを省略してs[a..]と書くとa文字目から末尾までの文字列スライスを取得できます。

したがって、s[n-3..]とすることで後ろ3文字の文字列スライスを取得することができます。
文字列の長さが3未満の時は当然teaと一致しないのですが、そのような文字列を評価しようとしてs[n-3..]とやってしまうとRustがpanicを引き起こしてしまうので、事前にNoを出力するように分岐しておきましょう。

use proconio::input;

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

    if n < 3 {
        println!("No");
        return;
    }

    if &s[n - 3..] == "tea" {
        println!("Yes")
    } else {
        println!("No")
    }
}

B - You're a teapot

$1\le |S| \le 100$と小さいため文字列中で両端がtであるような部分文字列を全探索し、その文字列内でいくつtが存在し、充填率はいくつかを$O(|S|^3)$で高速に求めることが可能です。

まず、両端がtであるような部分文字列を二重ループで求めます。
次に、左端を$i$、右端を$j$として、$i+1$から$j-1$までの範囲にtがいくつあるかを調べます。
ここで、この範囲内にtが存在しない場合は充填率は$0$なのでスキップします。tが存在する場合は充填率は正となるため、充填率を求めて充填率のmax値を更新します。

プログラム上では両端をはじめから考慮しないことにしているので、$x←x-2, \ |t|←|t|-2$としており、$|t|=(j-1)-(i+1)+1=j-i-1$なので、充填率は$x/(j-i-1)$として計算しています。

use proconio::input;
use proconio::marker::Chars;

fn main() {
    input! {
        s: Chars,
    }

    let mut ans = 0.0f64;
    for i in 0..s.len() {
        for j in (i + 1)..s.len() {
            if s[i] != 't' || s[j] != 't' {
                continue;
            }
            let mut x = 0.0f64;
            for k in i+1..j {
                if s[k] == 't' {
                    x += 1.0;
                }
            }
            if x == 0.0 {
                continue;
            }
            ans = ans.max(x / ((j - i - 1) as f64));
        }
    }
    println!("{}", ans);
}

C - Flush

この問題のように、一見どこから手をつければ良いかわからない問題では実験して糸口を掴むのが覿面です。

問題ページのサンプル1を観察してみましょう。

例）サンプル1
4 5
4 1 8 4
1
8
5
2
10

この問題において、$1\le B \le 17 = \sum A$の答えを表にまとめると以下のようになります。

$B$	1	2	3	4	5	6	7	8	9~17
$\min x$	1	5	8	11	14	15	16	17	-1

ここでまず着目していただきたい点は大きく二点です。

1. $B$が$1$から$2$へ、たった$1$大きくなっただけで答えが$4$も変わっていて、かつこの変化量は$N=4$の大きさと一致する
2. $9 \le B \le 17$においては答えは$-1$（勝利不可）で一定、かつ$9$は$\max{A_i}$よりも大きい最小の数である

このような観察結果をもとに考察を進めます。
偶然そうなっているだけということもありますが、面白い挙動をしている場合は往々にして答えの重要な糸口になることが多いものです。

まず、「$B$が$1$から$2$へ、たった$1$大きくなっただけで答えが$4$も変わっていて、かつこの変化量は$N=4$の大きさと一致する」の理由について考えてみましょう。
$B=2$のとき、ディーラーは、できるだけ一つのフレーバーが$2$個以上とならないように工夫してティーバッグを選ぼうとしますが、そのためには各フレーバーが一種類となるようできるだけ異なるフレーバーのティーバッグを選ぶことが最善の戦略となります。
したがって、$x=4$までは種類の異なるフレーバーを$1$つずつ選ぶことでディーラーが勝利します。
一方で、$x=5$のとき、$N=4$なので鳩の巣原理より必ずいずれかの種類のフレーバーは$2$個以上選ばなくてはならず、ディーラーは必ず負けてしまいます。
以上の考察で、「難易度$B$に対し、同じ種類のフレーバーが$B$個にならないよう可能な限り異なるフレーバーのティーバッグを選ぶ」 というのがディーラーの最善の戦略であることがわかりました。
その戦略に基づいているので、「$B$が$1$から$2$へ、たった$1$大きくなっただけで答えが$4$も変わって」いるのです。

なお、ここでもう一つ気づいておきたいのは、$B$が$2$から$3$へ変化した時の答えの変化量は$3$であり、$N=4$と一致しません。ここで、ディーラーは戦略 「難易度$B$に対し、同じ種類のフレーバーが$B$個にならないよう可能な限り異なるフレーバーのティーバッグを選ぶ」 に基づき行動するのですが、$A_2=1$のため2種類目のフレーバーは$2$個以上選択することができません。
つまり、各フレーバーで選択できる個数にはキャップ（上限）がある ということです。
したがって、ディーラーの最善の戦略は次のように再定義することができます：

ディーラーの最善の戦略は、 「難易度$B$に対して、各フレーバーの個数が$\min(B-1, A_i)$個以下となるように選ぶ」 ことである。（このような戦略のもとでフレーバーを選択できないとき、ディーラーの負け、すなわちあなたの勝ちとなる）

次に、「$9 \le B \le 17$においては答えは$-1$（勝利不可）で一定、かつ$9$は$\max{A_i}(i \in N)$よりも大きい最小の数である」の理由について考えてみましょう。
これは一つ目の理由と比べてかなり単純で、$B > \max A_i(i \in N)$のとき、当然どのフレーバーも個数が$\max A_i(i \in N)$以下なので、同じ種類のフレーバーを$B$個選択できないというだけです。例えば、サンプル1ではフレーバーの個数上限は$8$なので、難易度$B=9$は最初から実現不可能$=$あなたの負け確定であることがわかります。

以上でこの問題を解くための考察材料が出揃いました。
ディーラーの最善の戦略は、 「難易度$B$に対して、各フレーバーの個数が$\min(B-1, A_i)$個以下となるように選ぶ」 ことなので、ディーラーが勝利できる$x$の最大値は$\sum_{i=1}^N \min(B-1, A_i)$です。

サンプル1において、$B=4$での例を図で表すと以下のようになります。

このとき、ディーラーが勝利できる$x$の最大値は$\sum_{i=1}^N \min(B-1, A_i)=3+1+3+3=10$です。
したがって、あなたが勝利できる$x$の最小値は$11$となります。
また、$B>8$のとき、図より明らかにいずれかのフレーバーを$B$個以上選択することができないため、$x$は$-1$となります。

以上、プログラム上でどのように実装するかを考えましょう。
二分探索と累積和を用いて解くこともできますが、$B \le 10^6$のため$B=1 \cdots 10^6$までの答えを尺取り法を用いてあらかじめ全て求めておき、各クエリ$B_i$に対してはあらかじめ求めた答えを参照するだけという解法で解くこともできます。

$\min x$を$B=1 \cdots 10^6$について小さい順に求めていく上で、$\sum_{i=1}^N \min(B-1, A_i)$を高速に求めるため、配列$A$が昇順にソートされていると都合がいいです。
以下図を見ると一目瞭然ですが、$B$が大きくなるにつれて$\sum A'_i$で計算する水色の領域がだんだんと右に拡大していきます。($A'$は昇順ソートした$A$です)

ここで、水色の領域のうち最大の位置$j$と、水色領域に含まれる$A'_i$ の総和 $\text{sum} = \sum _{i=1}^j {A'_i}$を尺取り法というアルゴリズムで管理します。

そうすると、$B \le \max A$である限りあなたが勝利するような$x$が存在し、$x$の最小値は$\text{sum} + (n-j)\times (B-1)+1$のように表すことができます。$B > \max A$となったとき、それ以降の$B$については$x=-1$のままにしておきます。

あとは、クエリ$B_i$ごとに上記のアルゴリズムで前計算しておいた答えを参照するだけでよいです。
全体の計算量は$O(NlogN + \max B + Q)$となり、十分高速にこの問題を解くことができました。

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

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

    let mut a = a;
    a.sort();
    let mut map = vec![-1isize; 10usize.pow(6) + 1];
    let mut sum = 0;
    let mut j = 0;
    for i in 1..10usize.pow(6) + 1 {
        if a[n - 1] < i {
            break;
        }
        while j < n && a[j] < i {
            sum += a[j];
            j += 1;
        }
        map[i] = (sum + (n - j) * (i - 1) + 1) as isize;
    }

    let mut ans = vec![0isize; q];
    for i in 0..q {
        ans[i] = map[b[i]];
    }

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

D - XNOR Operation

解説では、非空の$S$に含まれる0の個数が奇数の時は必ず「美しい文字列」となるという性質を利用しています。

私はその性質に気づかずdp（動的計画法）で解いたのでその解法を紹介します。
なお、XNOR演算に結合則が成り立つ、すなわち$S$をどのような順序で「圧縮」（$|S|=1$となるまで、隣接する文字をXNOR演算の結果で置き換え）しても、最後に残る唯一の文字は変わらないという事実を利用しています。

$dp$テーブルを次のように定義します。
$$dp_{i,\text{bit}}:=左からi番目を右端とする部分文字列について、圧縮結果が\text{bit}となるような文字列の個数$$

このように定義すると、$dp_{i+1,0}$および$dp_{i+1,1}$への遷移式は次のように書くことができます。

• $S_{i+1}=$0のとき
  $dp_{i+1,0} = dp_{i, 1} + 1$ （10の圧縮結果が0となる、または$S_{i+1}$単体で0）
  $dp_{i+1,1} = dp_{i, 0}$ （00の圧縮結果が1となる）

• $S_{i+1}=$1のとき
  $dp_{i+1,0} = dp_{i, 0}$（01の圧縮結果が0となる）
  $dp_{i+1,1} = dp_{i, 1} + 1$（11の圧縮結果が1となる、または$S_{i+1}$単体で1）

そして、今回求めるべきは圧縮結果が1となる美しい文字列の総数なので、$dp_{-,1}$を$\text{sum}$に都度足し上げていくことにしましょう。
右端を$i$番目とするような文字列のうちで美しい文字列の個数を数え上げているため、美しい文字列が重複することはありません。

以上のようにして、$O(N)$で高速にこの問題を解くことができました。

補足

このような解法をどのようにして思いついたかを少しだけ書いておきます。
まず、相異なる$i,j$を選んでs[i..j]（$i$番目から$j-1$番目までの文字列スライス）が美しい文字列かどうか一つずつ判定する方法は、美しい文字列かどうかの判定を以下に高速にできたとしても$i,j$の全列挙で$O(N^2)$の計算量となってしまうため筋がいいとは言えません。

そこで、$i,j$どちらかを固定して、$i$を左端とするような部分文字列の個数か、$j$を右端とするような部分文字列の個数が差分更新かdpを利用してまとめて計算できないかを考えます。

また、XNOR演算の結合則により$S$はどのように圧縮しても結果が一意となるという事実もあらかじめわかっている必要があります。

そうすると、「仮に右端が$i$番目となるような部分文字列のうち圧縮結果が0または1となるような個数がわかっている場合、右端が$i+1$番目となるような部分文字列についても高速に求められないか」という思考が自然と生まれ、本記事で書いたような遷移を発見できればあとは実装に落とし込むだけ、となります。

他には、「バイナリ文字列の問題はdp（動的計画法）で解けることが多い」というのもヒントになるかもしれません。

use proconio::input;
use proconio::marker::Bytes;

fn main() {
    input! {
        n: usize,
        t: Bytes,
    }
    let mut dp = vec![vec![0usize; 2]; n];
    let mut sum = 0;
    if t[0] == b'1' {
        dp[0][1] = 1;
        sum = 1;
    } else {
        dp[0][0] = 1;
    }

    for i in 1..n {
        if t[i] == b'1' {
            sum += dp[i - 1][1] + 1;
            dp[i][1] += dp[i - 1][1] + 1;
            dp[i][0] += dp[i - 1][0];
        } else {
            sum += dp[i - 1][0];
            dp[i][1] += dp[i - 1][0];
            dp[i][0] += dp[i - 1][1] + 1;
        }
    }

    println!("{}", sum);
}

E - Trapezium

$N \le 2000$と$O(N^2)$解法でも十分に間に合うため、相異なる点$i,j$を選び、それらを結ぶ線分の傾きが同じであるようなペアの個数を求めれば良いです。
なお、$3$点が同一直線上にないことが問題文で保証されているため、傾きが同じ2線分（4頂点）を選ぶと直線になってしまうというケースは起こりえません。
実装としては[キー:傾き, 値:線分の個数]というハッシュマップを定義すればよいのですが、傾きを(dx/gcd(dx,dy), dy/gcd(dx,dy))のように$x$座標の差分と$y$座標の差分をそれぞれ$dx,dy$の最大公約数で割ったもののタプルで持っておくことで、$dy/dx$での除算誤差や、傾きが同じで長さが異なるだけの線分をカウントしない、といったつまらないミスを防ぐことができます。
このようにして、あるキーに対する値$\text{cnt}$について、$\text{cnt}\times (\text{cnt}-1)/2$を足し上げていけば良いだけ、と思いきや最後に一つだけ落とし穴があります。

それは、平行四辺形の場合です。
平行四辺形は一つにつき二つの平行な線分のペアを持っているため、ダブルカウントが生じます。
よって、最後に平行四辺形の個数を求め差し引いてあげましょう。
平行四辺形は、対辺の長さと傾きが等しい四角形であるため、(dx/gcd(dx,dy), dy/gcd(dx,dy), |dx|, |dy|)のような4要素を持つタプルを持つことで平行四辺形を成立させる線分ペアの個数をカウントすることができます。

以下の実装においては、少し変かもしれませんが台形についてはbox、平行四辺形についてはparallelというprefixをつけてそれぞれカウントしています。全体計算量は$O(N^2)$で十分高速にこの問題を解くことができました。

use proconio::input;

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

    let mut x = vec![0; n];
    let mut y = vec![0; n];
    for i in 0..n {
        input! {
            x_i: usize,
            y_i: usize,
        }
        x[i] = x_i;
        y[i] = y_i;
    }

    let mut parallel_map = std::collections::HashMap::new();
    let mut box_map = std::collections::HashMap::new();

    for i in 0..n {
        for j in (i + 1)..n {
            let dx = x[i] as i64 - x[j] as i64;
            let dy = y[i] as i64 - y[j] as i64;
            let g = gcd(dx.abs() as u64, dy.abs() as u64) as i64;
            let mut parallel_key = (dx / g, dy / g, dx.abs(), dy.abs());
            let mut box_key = (dx / g, dy / g);
            if parallel_key.0 < 0 || (parallel_key.0 == 0 && parallel_key.1 < 0) {
                parallel_key = (-parallel_key.0, -parallel_key.1, parallel_key.2, parallel_key.3);
                box_key = (-box_key.0, -box_key.1);
            }

            *parallel_map.entry(parallel_key).or_insert(0) += 1;
            *box_map.entry(box_key).or_insert(0) += 1;
        }
    }
    let mut box_cnt = 0i64;
    let mut parallel_cnt = 0i64;
    for &count in box_map.values() {
        box_cnt += count * (count - 1) / 2;
    }
    for &count in parallel_map.values() {
        parallel_cnt += count * (count - 1) / 2;
    }
    parallel_cnt /= 2;

    println!("{}", box_cnt - parallel_cnt);   
}
// gcd（最大公約数を求めるライブラリ）は省略

まとめ

今回のコンテストでは、Rustで初めてf64やpow()メソッド、Bytes型を使いました。
基礎的な道具については今後もD問題までを解いていきつつ覚えていけば良いのですが、E問題で利用したGCDライブラリがなかったりとライブラリ不足が深刻です。
いつかたっぷり時間をとってライブラリ整備をしたいですね、、、