累積和とは何か、およびABC098 C: Atttentionの解説

を解いていて、「累積和」という概念を知ったのでメモ。

こちらの本の「4.2 階差と累積和」でも紹介されている。

問題解決のための「アルゴリズム×数学」が基礎からしっかり身につく本

作者:米田優峻
技術評論社

Amazon

また、すでに

という素晴らしい記事があるので、ここでは累積和についての簡単な説明と、具体例としてABC098のC問題の解き方を紹介。

累積和とは

数列 a = [a_1, a_2, ... a_N] に対して

s_0 = 0  
s_1 = a_1
s_2 = a_1 + a_2
...
s_i = a_1 + a_2 + ... + a_i

つまりi番目までの和をとったものを累積和と呼ぶ。

累積和の何がうれしいのか

「配列の特定の区間の和を求める」ような処理を大量に行うシーンにおいて、
累積和を先に計算しておくことで、求めている計算結果を高速に得られる

ということだと理解した。

もう少し具体的に書くと、数列 a = [a_1, a_2, ... a_N] の l 番目から r 番目までの和を求める場合、O(N) の時間がかかる。

let sum = 0;
for (let i = l; i < r; i++) {
  sum += i;
}

これに対し、累積和を先に求めておくと、l 番目から r 番目までの和は

s[r+1] - s[l]

なので、O(1) で求まる。

なお、累積和の計算は、先頭の要素から順に足し合わせていけばいいので O(N) で求まる。

ABC098 C問題：Attentionを解く

ここで、具体例として冒頭に挙げた問題を考えてみる。

問題文は以下。

N 人の人が東西方向に一列に並んでいます。それぞれの人は、東または西を向いています。誰がどの方向を向いているかは長さ N の文字列 S によって与えられます。西から i 番目に並んでいる人は、S_i = E なら東を、S_i = W なら西を向いています。

あなたは、N 人のうち誰か 1 人をリーダーとして任命します。そして、リーダー以外の全員に、リーダーの方向を向くように命令します。このとき、リーダーはどちらの方向を向いていても構いません。

並んでいる人は、向く方向を変えるのを嫌っています。そのためあなたは、向く方向を変える人数が最小になるようにリーダーを選びたいです。向く方向を変える人数の最小値を求めてください。

まず、i 番目の人をリーダーに任命した場合、

1〜（i-1）番目までで「西」を向いている人
（i+1）〜N番目までで「東」を向いている人

は、それぞれ向く方向を変える必要がある。

これを、iを1番目からN番目までループしつつ計算すると、O(N²)の計算量になってしまう。

そこで、

s_i = 1番目からi番目までで、東（or 西）を向いている人の数の合計

と定義する。

そうすると、i番目の人がリーダーのときの求めたい人数は

1〜(i-1)番目までで「西」を向いている人
-> s_[i-1]


(i+1)〜N番目までで「東」を向いている人
-> (i+1)〜N番目までの人数 - (i+1)〜N番目までで「西」を向いている人
-> (N - i) - (s_N - s_[i])

両者を足し合わせることで導出できる。

以上を Rust で書くと、こんな感じ。

use proconio::{input, marker::Chars};

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

    let mut s = vec![];
    s.push(0);

    // s[i]: 1~i 番までで W を向いている人の数
    let mut count = 0;
    for &ch in &chars {
        if ch == 'W' {
            count += 1;
        }
        s.push(count);
    }

    // i 番目がリーダーだったとき、
    // 1~i-1 番目まで: W を向いている人が向く方向を変える => s[i-1] 人
    // i+1~n 番目まで: E を向いている人が向く方向を変える
    // => W を向いているのが s[n]-s[i] 人
    // => E を向いているのは {n-(i+1)+1} - (s[n]-s[i])
    let mut ans = n;
    for i in 1..=n {
        let target_num = s[i - 1] + n - i - (s[n] - s[i]);
        if ans > target_num {
            ans = target_num;
        }
    }
    println!("{}", ans);
}