RustでABC287復習

はじめに

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

復習

A - Majority

賛成(For)が過半数あるかを判定する問題。
Forの数え方は文字列Forをto_string()でStringに変換すると比較できるようになるので数えられる。
Forの数が$N/2$より大きいと過半数だと判定できる。

use proconio::input;


fn main() {
    input! {
        n: usize,
        s: [String; n],
    }
    let mut num_for = 0u64;
    for i in 0..n {
        if s[i] == "For".to_string() {
            num_for += 1;
        }
    }
    if n as u64 / 2 < num_for {
        println!("Yes");
    } else {
        println!("No");
    }
}

B - Postal Card

長さ6の文字列のうち、長さ3の文字列を末尾に持つものを数える問題。
全ての$S_i$の後ろ3文字を、全ての$T_j$と比較し、マッチしている$T_j$が一つ以上あれば答えを1つ増やす。
$S_i$と$T_j$がマッチしているかとマッチしている$T_j$があるかはフラグで管理すると答えが出せる。

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


fn main() {
    input! {
        n: usize,
        m: usize,
        s: [Chars; n],
        t: [Chars; m],
    }
    let mut ans = 0u64;
    for i in 0..n {
        let mut is_valid = false;
        for j in 0..m {
            let mut is_same = true;
            for k in 0..3 {
                if s[i][3+k] != t[j][k] {
                    is_same = false;
                }
            }
            if is_same {
                is_valid = true;
            }
        }
        if is_valid {
            ans += 1;
        }
    }
    println!("{}", ans);
}

C - Path Graph?

与えられたグラフがパスグラフか判定する問題。
パスグラフとは一本道みたいな形をしているグラフのことを指すらしい(参考)。
グラフが一本道であるためには以下の3つの条件を満たす必要がある。

• 辺の数は$N-1$
• ある辺から他の全ての辺に行くことが可能
• 各グラフは2より多くの辺を持たない

上の二つの条件はこのパスグラフが木になっているからである。なぜなら、このグラフは連結で、閉路を持たない単純(無向)グラフだからである。
一番下の条件は木が一直線になるために必要な条件である。全ての頂点から辺が2つずつ出ている場合はパスグラフにならないが木である以上そのようなことは起こりえない。
ある辺から全ての辺に行くことが可能かの判定は適当な頂点を根として幅優先探索をすることで行う。
辺の数や各頂点の次数(頂点から出ている辺の数)は$M$の数や幅優先探索の過程で到達した頂点で適宜調べると求められる。
以上で全ての条件の判定ができるので答えが求まる。

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


fn main() {
    input! {
        n: usize,
        m: usize,
        edges: [(Usize1, Usize1); m],
    }
    if m != n - 1 {
        println!("No");
        return;
    }
    let mut graph = vec![vec![]; n];
    for &(u, v) in &edges {
        graph[u].push(v);
        graph[v].push(u);
    }
    let mut is_visited = vec![false; n];
    let mut stack = vec![0usize];
    while let Some(node) = stack.pop() {
        if is_visited[node] {
            continue;
        }
        is_visited[node] = true;
        if 2 < graph[node].len() {
            println!("No");
            return;
        }
        for &next_node in &graph[node] {
            stack.push(next_node);
        }
    }
    let mut is_valid = true;
    for i in 0..n {
        if !is_visited[i] {
            is_valid = false;
        }
    }
    if is_valid {
        println!("Yes");
    } else {
        println!("No");
    }
}

D - Match or Not

文字列$S$の前と後ろを組み合わせて文字列$T$が作れるかを判定する問題。
まず、$S$と$T$の先頭から$i+1$文字目までが合っているかを$\mathrm{is\_match\_front}[i]$で、後ろから$|T|$文字目から$i+1$文字目までが合っているかを$\mathrm{is\_match\_back}[i]$で管理する。これは先頭と後ろから舐めていき、一致しないものが出てくるまで$\mathrm{true}$を入れ続けると作ることができる。
あとは各$x$に対して

• $x=0$なら$\mathrm{is\_match\_back}[0]$
• $x = |T|$なら$\mathrm{is\_match\_front}[|T|-1]$
• それ以外なら$\mathrm{is\_match\_front}[x-1]\quad \& \quad\mathrm{is\_match\_back}[x]$

を判定すれば答えが求められる。

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


fn main() {
    input! {
        s: Chars,
        t: Chars,
    }
    let n = s.len();
    let m = t.len();
    let mut is_match_front = vec![false; m];
    for i in 0..m {
        if s[i] == t[i] || s[i] == '?' || t[i] == '?' {
            is_match_front[i] = true;
        } else {
            break;
        }
    }
    let mut is_match_back = vec![false; m];
    for i in 0..m {
        if s[n-i-1] == t[m-i-1] || s[n-i-1] == '?' || t[m-i-1] == '?' {
            is_match_back[m-i-1] = true;
        } else {
            break;
        }
    }
    for i in 0..=m {
        println!("{}", if i == 0 {
            if is_match_back[0] {
                "Yes"
            } else {
                "No"
            }
        } else if i == m {
            if is_match_front[m-1] {
                "Yes"
            } else {
                "No"
            }
        } else {
            if is_match_front[i-1] && is_match_back[i] {
                "Yes"
            } else {
                "No"
            }
        });
    }
}

E - Karuta

文字列の決まり字が何文字かを求める問題。
ABC283Gと似た要領で解くことができる。
先頭から$i$文字目まで同じものでグループに分けていき、グループの中の要素数が1以下になるまで再帰的に繰り返す。このときの$i-1$が答えになる。
計算量はアルファベット26種類なので$O(26^N)$になりそうに感じるが実際は各$S_i$の各文字をたかだか一度しか見ないので$O(\sum_{i=1}^n|S_i|)$になるため十分間に合う。

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


fn recursive(i: usize, indexes: &Vec<usize>, s: &Vec<Vec<char>>, mut ans: &mut Vec<u64>) {
    if indexes.len() < 2 {
        for j in 0..indexes.len() {
            ans[indexes[j]] = i as u64 - 1;
        }
        return;
    }
    let mut initials = vec![vec![]; 26];
    for j in 0..indexes.len() {
        if s[indexes[j]].len() <= i {
            ans[indexes[j]] = i as u64;
            continue;
        }
        initials[(s[indexes[j]][i] as u8 - 'a' as u8) as usize].push(indexes[j]);
    }
    for j in 0..26 {
        recursive(i+1, &initials[j], &s, &mut ans);
    }
}

fn main() {
    input! {
        n: usize,
        s: [Chars; n],
    }
    let indexes = (0..n).collect::<Vec<usize>>();
    let mut ans = vec![0u64; n];
    recursive(0, &indexes, &s, &mut ans);
    for &a in &ans {
        println!("{}", a);
    }
}

F - Components

木から頂点を消したときの連結成分の数ごとのパターンを求める問題。
$N = 5000$で$1 \le i \le N$に対して答えを出す必要があるので一つ当たり$O(N)$程度で出す必要がある。
グラフが木であることを踏まえると葉から根にかけて各頂点で$N$種類の答えを更新していくとよさそうに見える。
更新の方法は今見ている頂点を根とした部分木を考えて、その部分木で連結成分数が$i$になる場合の数を$O(1)$で更新する方法を考えると良い。ただし、そのような方法は(少なくとも私は知ら)ないので別の方法を考える。
そうすると、実は全ての頂点で親と子でそれぞれの連結成分数のパターンを考慮しながら更新しても間に合う。
具体的には、現在見ている部分木の頂点を使用する場合($s=1$)と使用しない場合($s=0$)に連結成分数が$i$になる場合の数を$\mathrm{dp}[s][i]$、子の頂点も同様にその頂点を根としたときの部分木を見たときの連結成分数が$j$になる場合の数を$\mathrm{next\_dp}[s][j]$と置くと

dp[0][k] = \sum_{i=0}^k \mathrm{dp}[0][i] \times (\mathrm{next\_dp}[0][k-i] + \mathrm{next\_dp}[1][k-i]) \\
dp[1][k] = \sum_{i=0}^k \mathrm{dp}[1][i] \times \mathrm{next\_dp}[0][k-i] + \sum_{i=0}^{k+1} \mathrm{dp}[1][i] \times \mathrm{next\_dp}[1][k-i-1]

のように遷移する。
気持ちとしては下の図のように基本は連結数は合計した数になるが、両方の頂点を使用する場合だけ根の連結成分同士が連結するため1つ減る。

この遷移は各辺がつなぐ木の親と子を根とする部分木の成分数の最大値の二重ループになるため計算量は$O(N^3)$になりそうに見える。
しかし、実際は各部分木で連結成分数としてありえる最大値同士の二重ループであればならし計算量が$O(N^2)$になるため間に合う。
このようになる理由を説明する。
まず連結成分数の最大値はその部分木に含まれる頂点の数$i$である。(実際はもっと少ない)
$i=1$のとき、$i \times i = i^2$より成立する。
連結成分数が$i$と$j$の部分木で遷移を行う時ループの回数は
$i^2 + j^2 + ij \lt i^2 + 2ij + j^2 = (i+j)^2$より成立する。
よってこの方法で間に合い、答えが出る。

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


const MOD: u64 = 998244353;

fn dfs(src: Option<usize>, node: usize, graph: &Vec<Vec<usize>>) -> Vec<Vec<u64>> {
    let mut max_num = 1usize;
    let mut dp = vec![vec![0u64; max_num+1]; 2];
    dp[0][0] = 1u64;
    dp[1][1] = 1u64;
    for &next_node in &graph[node] {
        if let Some(s) = src {
            if next_node == s {
                continue;
            }
        }
        let next_dp = dfs(Some(node), next_node, &graph);
        let num = next_dp[0].len() - 1;
        let mut new_dp = vec![vec![0u64; max_num+num+1]; 2];
        for i in 0..=max_num {
            for j in 0..=num {
                new_dp[0][i+j] = (new_dp[0][i+j] + (dp[0][i] * (next_dp[0][j] + next_dp[1][j])) % MOD) % MOD;
                new_dp[1][i+j] = (new_dp[1][i+j] + dp[1][i] * next_dp[0][j]) % MOD;
                if 0 < i + j {
                    new_dp[1][i+j-1] = (new_dp[1][i+j-1] + (dp[1][i] * next_dp[1][j])) % MOD;
                }
            }
        }
        max_num += num;
        dp = new_dp;
    }
    return dp;
}

fn main() {
    input! {
        n: usize,
        edges: [(Usize1, Usize1); n-1],
    }
    let mut graph = vec![vec![]; n];
    for &(u, v) in &edges {
        graph[u].push(v);
        graph[v].push(u);
    }
    let ans = dfs(None, 0usize, &graph);
    for i in 1..=n {
        println!("{}", (ans[0][i]+ans[1][i]) % MOD);
    }
}