ABC221 を Rust で解く (E問題まで)

はじめに

AtCoder Beginner Contest 221 の E問題までを Rust で解きました！
詳しめに解説をしていこうと思います。
役にたちましたら LGTM いただけると嬉しいです。
AC 確認はしてますが、もっとスマートな書き方がありましたら教えてください。

なぜやるのか

コンテストから時間をあけて解きなおすことで、記憶へ定着すればいいなと思っています。
また、会社で Rust を使うのですが私はまだ慣れておらず、その練習のためにやっています。
(コンテスト本番には C++ を使って参加している水色コーダーです。)

目次

A - Seismic magnitude scales (Difficulty 10)
B - typo (Difficulty 58)
C - Select Mul (Difficulty 379)
D - Online games (Difficulty 832)
E - LEQ (Difficulty 1515)

前提

この記事にあるコードはこちらを import しています。

#[allow(unused_imports)]
use proconio::{input, fastout, marker::Chars};
#[allow(unused_imports)]
use std::collections::{HashSet, HashMap, BTreeSet, VecDeque};
#[allow(unused_imports)]
use std::cmp::{max, min};
#[allow(unused_imports)]
use itertools::Itertools;

A - Seismic magnitude scales

a - b の回数だけ 32 をかければ答えになります。
制約より、オーバーフローの心配もありません。
( $32 = 2^5$ を最大で6回かけるので、答えは最大で $2^{30}$ になります)

ちなみに、実際のマグニチュードは1増えるごとにエネルギーが $\sqrt{1000}$ 倍になります。32倍ではないのでご注意を。

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

  let mut ans = 1;
  for _ in 1..=a-b {
    ans *= 32;
  }
  println!("{}", ans);
}

B - typo

1. s の i 文字目と i + 1 文字目を入れ替える
2. s と t を比較する
3. s を元に戻す

というアルゴリズムで解きました。
for の範囲と、はじめから等しいときに注意です。

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

  if s == t {
    println!("Yes");
    return;
  }

  for i in 0..s.len() - 1 {
    s.swap(i, i + 1);
    if s == t {
      println!("Yes");
      return;
    }
    s.swap(i, i + 1);
  }
  println!("No");
}

C - Select Mul

全探索の問題ですね。

N は1文字ずつ処理したいので、数値ではなく文字列として扱います(Rust では Chars)。
2つの Vec を用意して、 N を1文字ずつどちらかに振り分けます。
1文字あたり2通りの振り分け先があるので、全体で $2^{|N|}$ 通りを試すことになります。
こういうときは bit 全探索を使いましょう。

振り分けた後は、降順にソートして数値に変換します。こうすることで数値が最大化されます。
あとは、変換結果をかけた数値を比較して、そのなかの最大値が答えになります。

振り分け方が $2^{|N|}$ 通りで、ソートに $O(|N| log |N|)$ かかるので、全体の計算量は $O(2^{|N|} * |N| log |N|)$ になります。

fn main() {
  input! {
    n: Chars
  }

  let mut ans = 0;
  for mask in 0..1 << n.len() {
    let mut s = vec![];
    let mut t = vec![];
    for i in 0..n.len()
    {
      if (mask >> i & 1) == 1 { s.push(n[i]) }
      else { t.push(n[i]) }
    }
    s.sort();
    t.sort();
    let s = s.iter().rev().collect::<String>().parse::<usize>().unwrap_or(0);
    let t = t.iter().rev().collect::<String>().parse::<usize>().unwrap_or(0);
    ans = max(ans, s * t);
  }
  println!("{}", ans);
}

D - Online games

イベントソートの問題ですね。

ABC118-D と似ています。Difficulty も近いですね。再登場ということで今回のほうが Difficulty が低いです。

一見するといもす法で解けそうなのですが、計算量が $O(N + max(A_i + B_i))$ となり、今回の制約から解けません。
いもす法を使うには座標圧縮などの工夫が必要です。
それでやってもよかったのですが、実装が相対的に複雑になるので、今回はイベントソートを使います。

イベントソートは、各プレイヤーがログインしてからログアウトするまでをそれぞれイベントとして扱います。
ログインした日を +1、ログアウトした日を -1 として、この値の累積値を追います。ここはいもす法と同じですね。
また、ひとつ前のイベントを記憶しておいて、今回と前回の差が、求める期間になります。

このイベントを平衡二分探索木で管理して、0日目から順に処理することで $O(N log N)$ で解くことができました。

(この問題では人数が0人だった日は不問ですね。もしこれも答える場合は、0日目と最終日にイベントを追加すれば解けます(いわゆる番兵というやつです)。)

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

  let mut event: HashMap<usize, i64> = HashMap::new();

  for (a, b) in ab.into_iter() {
    *event.entry(a).or_insert(0) += 1; // ログインした日
    *event.entry(a + b).or_insert(0) -= 1; // ログアウトした日
  }

  let mut event: Vec<(usize, i64)> = event.into_iter().collect();
  event.sort();
  
  let mut prev = 0;
  let mut now_val = 0;
  let mut ans = vec![0; n + 1];
  for (key, val) in event.into_iter() {
    ans[now_val as usize] += key - prev; // [key - prev]日間が now_val 人だった
    now_val += val;
    prev = key;
  }
  for i in 1..=n {
    print!("{}", ans[i]);
    print!("{}", if i == n { '\n' } else { ' ' });
  }
}

E - LEQ

組み合わせの個数を数える問題ですね。誤読しやすかったような、、(私は誤読しました。)

まず、数字を2つ選んで、$A_i, A_j (i < j) $ とします。この2数でペアを作った時の組み合わせは $(2^{j - i - 1})$ になりますね。これが基本になります。
この $i, j$ を全探索すれば簡単に答えがわかるのですが、そうすると $O(N^2)$ かかってしまうので時間内に解けません。高速化しましょう。

数字の大小を扱うので、ソートすることを考えてみます。これは ABC216-F と似ていますね。ある数について考えると、それ以下の数字について考えれば大丈夫です。つまり、昇順にソートして前から見ていけばいいですね。ただ、今回の問題ではソート前の位置も必要なので注意です。これは enumerate() がつかえます。

続いて、ある数とペアになれる数について考えます。その中でも、ソート前の位置が自分よりも前にあるものを数え上げの対象にします。これは、前と後ろを対象にすると二重で数えられてしまう組が出てくるので、それを回避するためです。

いきなり実例ですが、ある数を $A_4$ について考えます。$A_4$ よりも前にある $A_0,A_1,A_2,A_3$ のなかで、 $A_0, A_2$ が $A_4$ 以下だとします。すると、求める組み合わせの個数は $A_0, A_4$ の組と $A_2, A_4$ について考えればいいので、それぞれ計算して $2^3 + 2^1$ になりますね。

これと、上までの流れを合わせて考えます。上の実例では、$A_0,A_2$ が計算済みで、$A_1, A_3$ は未計算です。これを言い換えると、計算済みの数字に対して総和を計算したことになります。さらに言い換えると、$A_0$ から $A_3$ までの区間和を計算したことになります。
区間和といえば累積和ですが、値が更新されるので今回は使えません。

1点更新・区間和を高速に計算できるものはありますでしょうか。あります。
データ構造の SegmentTree か FenwickTree(Binary Indexed Tree: BIT) が使えます。1点更新と区間和をどちらも $O(log N)$ で計算できます。全体計算量が $O(N log N)$ になるので今回は大丈夫そうです。また、実行速度と実装の簡単さを考えて、私は FenwickTree を使いました。

さて、ここまでで FenwickTree を使えば解けそうなことがわかりました。

ですが、もう一段階あります。区間和を求めるときに、自分の位置によって計算対象の値が変わってしまいます。
$A_1$ が $A_0$ とペアを作るときは $2^0$ 通り、$A_2$ が $A_0$ とペアを作るときは $2^1$ 通り、、、$A_j$ が $A_0$ とペアを作るときは $2^{j - 1}$ 通り。
となって、同じ $A_0$ でも距離によって値が変わってしまいます。

ここでは、距離を値に変換します。といってもやることは簡単です。FenwickTree には $2^{-(index)-1}$ を登録して、総和に対して $2^{(index)}$ をかけた値を数え上げに使います。$2^{(index)}$ をかけることで、各数について自分の位置からみたペアの数を計算できます。このあたりは、指数に注目しながら紙に書くとわかりやすいかなと思います。

また、$2^{-1}$ は $2$ の逆元であり、mod を取る関係上、フェルマーの小定理より、998244353 - 2 乗をすれば求まります。この値を i 乗すれば、$2^{-i}$ の値も求められます。

さて、実装です。
FenwickTree について調べると Crate fenwick_tree があったので使いたかったのですが、AtCoder では使えませんでした。言語アップデートに期待です。というわけで、ここで扱うFenwickTree は自分で書くことにします。こちらの記事 より拝借しました(型を i64 にして mod を計算するなど、いくらか変えています)。

// a^n % m を計算する O(log n)
fn powmod(a: i64, n: i64, m: i64) -> i64 {
  let mut ret = 1;
  let mut a = a;
  let mut n = n;
  while n > 0 {
    if n % 2 == 1 { ret =  ret * a % m }
    a = a * a % m;
    n /= 2;
  }
  ret
}

#[inline]
fn lsb(i: usize) -> usize {
    i & i.wrapping_neg()
}

struct Fenwick {
  table: Vec<i64>,
}
impl Fenwick {
  pub fn new(size: usize) -> Self {
      Self { table: vec![0; size] }
  }
  pub fn prefix_sum(&self, mut i: usize, m: i64) -> i64 {
    let mut res = 0;
    while i != 0 {
        res += self.table[i];
        res %= m;
        i -= lsb(i);
    }
    res
  }
  pub fn add(&mut self, mut i: usize, x: i64, m: i64) {
    i += 1;
    while i < self.table.len() {
        self.table[i] += x;
        self.table[i] %= m;
        i += lsb(i);
    }
  }
}

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

  let mut a = a
    .into_iter()
    .enumerate()
    .collect_vec();
  
  // val が等しい場合は idx の小さいものから見る必要があるので、 sort_unstable だと WA になる
  a.sort_by_key(|v| v.1);

  let m: i64 = 998244353;
  let inv2 = powmod(2, m - 2, m);

  let mut fw = Fenwick::new(n);
  let mut ans: i64 = 0;
  for (idx, _) in a.into_iter() {
    ans += (fw.prefix_sum(idx, m) * powmod(2, idx as i64, m)) % m;
    ans %= m;
    fw.add(idx, powmod(inv2, idx as i64 + 1, m), m);
  }
  println!("{}", ans);
}

---

最後まで読んでいただきありがとうございます！