RustでABC284復習

はじめに

Rustの勉強のためにAtCoderのコンテストでRustを使うことにしました。
解説を読んだりして自分がわかった範囲で復習します。

復習

A - Sequence of Strings

愚直に実装する。
rev()関数を使うことで逆順でイテレートできる。

use proconio::input;


fn main() {
    input! {
        n: usize,
        s: [String; n],
    }
    for i in (0..n).rev() {
        println!("{}", s[i]);
    }
}

B - Multi Test Cases

愚直に実装する。
計算量は$O(TN)$で$1 \le TN \le 10^4$なため十分間に合う。

use proconio::input;


fn main() {
    input! {
        t: usize,
    }
    for _ in 0..t {
        input! {
            n: usize,
            a: [u64; n],
        }
        let mut ans = 0u64;
        for &ai in &a {
            if ai % 2 == 1 {
                ans += 1;
            }
        }
        println!("{}", ans);
    }
}

C - Count Connected Components

連結成分の数を求める問題なのでUnionFindを使う。
辺が結ぶ2頂点をuniteして同じグループにする。
グループの数が連結成分の数なのでルートの数を数えると答えになる。

use proconio::input;
use proconio::marker::Usize1;


struct UnionFind {
    roots: Vec<i64>,
}

impl UnionFind {
    fn new(n: usize) -> Self {
        UnionFind {
            roots: vec![-1; n],
        }
    }

    fn root(&mut self, x: usize) -> usize {
        if self.roots[x] < 0 {
            return x;
        } else {
            self.roots[x] = self.root(self.roots[x] as usize) as i64;
            return self.roots[x] as usize;
        }
    }

    fn is_same(&mut self, x: usize, y: usize) -> bool {
        self.root(x) == self.root(y)
    }

    fn size(&mut self, x: usize) -> usize {
        let root_x = self.root(x);
        (-self.roots[root_x]) as usize
    }

    fn unite(&mut self, x: usize, y: usize) -> bool {
        if self.is_same(x, y) {
            false
        } else {
            let mut root_x = self.root(x);
            let mut root_y = self.root(y);
            if self.size(root_x) < self.size(root_y) {
                std::mem::swap(&mut root_x, &mut root_y);
            }
            self.roots[root_x] -= (self.size(root_y)) as i64;
            self.roots[root_y] = root_x as i64;
            true
        }
    }
}

fn main() {
    input! {
        n: usize,
        m: usize,
        edges: [(Usize1, Usize1); m],
    }
    let mut uf = UnionFind::new(n);
    for &(u, v) in &edges {
        uf.unite(u, v);
    }
    let mut ans = 0u64;
    for i in 0..n {
        if i == uf.root(i) {
            ans += 1;
        }
    }
    println!("{}", ans);
}

D - Happy New Year 2023

N = p^2 q

となる素数$p$、$q$を求める問題。
$T$個のテストケースに対して行う必要はあるが、$T$は10以下なのであまり気にしなくてもいい。
$p$、$q$の小さい方を仮に$r$とすると

8 \le r^3 \lt p^2 q \le 9 \times 10^{18} \qquad (2 \le r)\\
r < 10^6

となるので$r$を$2$から順番に$N$を割り切ることができるか試しながら探す。このとき小さい方から探すと$p$、$q$が素数なので先に合成数が見つかることはない。
次に見つかった$r$が$p$と$q$のどちらなのかを調べる。
これは簡単で$N$が$r^2$で割り切れるかを確認すると良い。割り切れると$p = r$で、割り切れないと$q = r$になる。
$p = r$だった場合、$q = \frac{N}{r^2}$になり、$q = r$だった場合$p = \sqrt{\frac{N}{r}}$になる。
Rustで整数型の平方根の求め方はわからなかったので今回は二分探索で求めた。
これで計算量は全体で$O(T(\sqrt[3]{N} + \log \sqrt[\frac{2}{3}]N))$となり間に合う。

use proconio::input;


fn main() {
    input! {
        t: usize,
    }
    for _ in 0..t {
        input! {
            n: u128,
        }
        let mut div = 2u128;
        while 0 < n % div {
            div += 1;
        }
        let (p, q) = if n % (div * div) == 0 {
            (div, n / (div * div))
        } else {
            let mut ok = 1u128;
            let mut ng = n / div + 1;
            let n = n / div;
            while 1 < ng - ok {
                let target = (ok + ng) / 2;
                if (target * target) <= n {
                    ok = target;
                } else {
                    ng = target;
                }
            }
            (ok, div)
        };
        println!("{} {}", p, q);
    }
}

E - Count Simple Paths

頂点1を始点とするユニークなパスを数える問題。
全探索すると計算量が$O(N!)$になって間に合わないので全ての頂点の次数が$10$以下になることを活かして計算したくなりますが全探索します。
答えはユニークなパスの数か$10^6$の小さい方になるので、一旦全探索していき$10^6$個見つかった時点で全探索を終了すると十分間に合う。
パスの見つけ方は深さ優先探索(DFS)を使用する。
頂点1から始めて、訪れた頂点をセットで持つ。再帰関数から抜けるときにセットから自分の頂点を削除すると全探索できる。

use proconio::input;
use proconio::marker::Usize1;
use std::collections::HashSet;


const LIMIT: u64 = 1_000_000;

fn dfs(node: usize, num: u64, graph: &Vec<Vec<usize>>, mut visited: &mut HashSet<usize>) -> u64 {
    if !visited.insert(node) {
        return num;
    }
    let mut num = num+1;
    if LIMIT < num {
        return LIMIT;
    }
    for &next_node in &graph[node] {
        num = dfs(next_node, num, &graph, &mut visited);
    }
    visited.remove(&node);
    return num;
}

fn main() {
    input! {
        n: usize,
        m: usize,
        edges: [(Usize1, Usize1); m],
    }
    let mut graph = vec![vec![]; n];
    for &(u, v) in &edges {
        graph[u].push(v);
        graph[v].push(u);
    }
    let mut visited = HashSet::new();
    println!("{}", dfs(0usize, 0u64, &graph, &mut visited));
}

F - ABCBAC

文字列の中に分割された文字列$S$と分割されていない逆順の文字列$S$があるかを求める問題。
このような文字列を探す問題はZアルゴリズムが使えるように先頭に探したい文字列を持ってきて全体のどこに同じ文字列があるかを探す場合が多い気がしていて、今回もそのパターン。(コンテスト時間内に並び替えの方法に気づけなかったが・・・)
ここで文字列$U$を逆順にした文字列を$\bar{U}$として表記する。
また、分割された$S$の前部を$A$、後部を$B$とおく。
そうすると

S = AB

であり、$f_i(S)=T$でとなっている場合、

T = A \bar{B} \bar{A} B

となっている。ここで文字列$AB$は文字列$A$と文字列$B$がこの順番でつながっていることを示す。
また、文字列$U$の長さを$|U|$と表現するとき、

|AB| = N \\
|A| = |\bar{A}| \\
|B| = |\bar{B}| \\

が成り立つので、

|A \bar{B}| = |\bar{A} B| = N

が成り立つ。
なので$T$を前半と後半に分けることで$A\bar{B}$と$\bar{A}B$にわけることができる。
ここで

C = A\bar{B} \\
D = \bar{A}B

とおくと

C\bar{D} = A\bar{B}\bar{B}A \\
\bar{C}D = B\bar{A}\bar{A}B

となる。
ここでこの二つの文字列にZアルゴリズムを使用することで$A$と$B$を求めることができる。
Zアルゴリズムは連続している部分文字列が全体の文字列$U$の先頭何文字と同じかを$O(|U|)$で求めるものである。(ここがわかりやすかった)
なので2つの文字列からZアルゴリズムで求めた共通の文字数の和が$N$となる場所を求めると答えになる。見つからないとそのような$S$と$i$の組が存在しない。

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


fn z_algorithm(s: &Vec<char>) -> Vec<u64> {
    let n = s.len();
    let mut z = vec![0u64; n];
    z[0] = n as u64;
    let mut i = 1usize;
    let mut j = 0usize;
    while i < n {
        while i + j < n && s[i+j] == s[j] {
            j += 1;
        }
        z[i] = j as u64;
        if j == 0 {
            i += 1;
            continue;
        }
        let mut k = 1usize;
        while k < j && k + (z[k] as usize) < j {
            z[i+k] = z[k];
            k += 1;
        }
        i += k;
        j -= k;
    }
    return z;
}

fn main() {
    input! {
        n: usize,
        t: Chars,
    }
    let l = 2 * n;
    let mut rev_t = t.clone();
    rev_t.reverse();
    let mut s1r2r21 = vec![];
    s1r2r21.extend_from_slice(&t[0..n]);
    s1r2r21.extend_from_slice(&rev_t[0..n]);
    let z1 = z_algorithm(&s1r2r21);
    let mut s2r1r12 = vec![];
    s2r1r12.extend_from_slice(&rev_t[n..l]);
    s2r1r12.extend_from_slice(&t[n..l]);
    let mut z2 = z_algorithm(&s2r1r12);
    z2.push(0);
    for i in 0..n {
        if z1[l-i-1] + z2[l-n+i+1] != n as u64 {
            continue;
        }
        let mut ans = vec![];
        for j in 0..=i {
            ans.push(t[j]);
        }
        for j in i+n+1..l {
            ans.push(t[j]);
        }
        println!("{}", ans.iter().collect::<String>());
        println!("{}", i+1);
        return;
    }
    println!("-1");
}

G - Only Once

この問題で数列$A$を一つ定めると各頂点$k$が頂点$A_{B_{i, k}}$へ有向な辺を持つグラフができる。

「数列$B_i$の中でちょうど1度だけ出てくる数の種類数」は「頂点$i$から到達できる頂点のうち、閉路に含まれていない頂点の数」になる。
なぜなら、

• $B_i$に出てくる数は頂点$i$を含めて頂点$i$から到達できる頂点の番号
• $10^{100}$は十分大きいので($2N$以上なので)閉路に含まれる頂点は2回以上出現
  
  また、頂点$i$から到達可能な頂点の数が頂点$i$を含めて$l$個あるとき、「頂点$i$から到達できる頂点のうち、閉路に含まれていない頂点の数」は最後の頂点の接続先を変更することで$0$(最後の頂点が$i$に接続するとき)から$l-1$(最後の頂点が自分の頂点に接続するとき)までの数を取る。
  
  なので、$1$以上$N$以下の$l$に対して、到達可能な頂点の数が$l$となる数列$A$のパターン数を数え、それぞれについて$\frac{l(l-1)}{2}$で1度だけ出てくる数の種類数を数える。
  到達可能な頂点の数が$l$となる数列のパターン数は
• 頂点$i$
• 頂点$i$から到達可能な頂点
• 頂点$i$から到達不可能な頂点
  の3つについて考えるとよい。
  頂点$i$の選び方は$N$通りある。
  頂点$i$から到達可能な頂点のパターンは頂点$i$を固定するとそこから順番を考えた$l-1$個の頂点の選び方なので、頂点$i$の選び方と合わせて$ _NP_l$通りになる。
  頂点$i$から到達不可能な頂点のパターンは、$N-l$個のそれらの頂点は何になっていても問題ないことから$N^{N-l}$通りになる。
  これらは事前に$ _NP_l$と$N^{N-l}$を計算しておくことで$O(1)$で計算できる。
  そのため、これを$N$個(実際は$N-1$が最大なので$N-1$個)の$l$に対して行うことで計算量$O(N)$でこの問題が解ける。

use proconio::input;


fn main() {
    input! {
        n: u64,
        m: u64,
    }
    let mut ans = 0u64;
    let mut fact =vec![1u64; n as usize + 1];
    let mut exp = vec![1u64; n as usize + 1];
    for i in 1..=n {
        fact[i as usize] = ((n - i + 1) * fact[i as usize - 1]) % m;
        exp[i as usize] = (n * exp[i as usize - 1]) % m;
    }
    for l in 2..=n {
        let pattern = (((fact[l as usize] * exp[(n-l) as usize]) % m) * ((l * (l-1) / 2) % m)) % m;
        ans = (ans + pattern) % m;
    }
    println!("{}", ans);
}