Rustコトハジメ

プログラミング言語Rustと競プロに関する情報をお届けします。

しゃくとり法を図で考える

競プロで頻出のしゃくとり法。

例題としては、

があります。

しかし、しゃくとり法でなぜ正しいのか?なぜオーダーがO(N)なのか。しゃくとり法を実装するにしても、なかなかきれいに実装出来ないという人は少なくないはず。というか自分がそうである。

そこで、しゃくとり法について、その原理と実装方法について考えてみることにしました。

原理

しゃくとり法は、i,jが与えられたら計算は出来るが、ナイーブにやるとO(N2)かかる問題をO(N)で計算する手法です。この超のつく高速化のために強い制約が必要です。

しゃくとり法の直感的な理解としては、iを固定した時に限界までjを進めて、jを進めていってだめになったら次はiを進めて・・・を繰り返すということですが、なぜこれが出来るかというと、

f:id:akiradeveloper529:20190828122018j:plain

  • jを進めていくといつかFにぶつかる
  • iを進めていくといつかTにぶつかる

という性質があるからです。この性質が保証されていることが、しゃくとり法が実行可能な制約といえます。もちろん、i<=jです。

この図を、赤線のように辿ることによって、O(2N)の計算量になるということです。

なぜ他のところは見なくてもよいのでしょうか?

この図で、左下を(0,0)として説明します。すると、

  • (1,0)を見なくていい理由は、i<=jを満たさないから
  • (1,4)を見なくてよい理由というのは、どうせFだから
  • (1,1)を見なくてよい理由は、問題の制約による。見ないといけない問題だとしゃくとり法は使えないはず

とわかります。他の部分も同様です。この議論は、iとjを反転させたり、TとFを反転させても成立します。

実装

では、しゃくとり法をどのように実装するのが最善でしょうか?

D - Enough Array

を例にして考えます。

これはサンプル1でいうと、

f:id:akiradeveloper529:20190828135014j:plain

しゃくとり法をして、Tの数を数え上げろという問題と考えることが出来ます。そしてこの例でいうと、赤線のように、際のTを辿っていって、それより上にあるTを数え上げるプログラムを書けばよいとわかります。

fn solve() {
    input!{
        N:usize,K:usize,
        A:[usize;N],
    }
    let mut ans = 0;
    let mut acc = 0;
    let mut i = 0;
    let mut j = 0;
    for i in 0..N {
        while acc < K {
            if j==N {
                break;
            }
            acc += A[j];
            j += 1;
        }
        if acc >= K {
            ans += 1+(N-j);
        }
        acc -= A[i];
    }
    println!("{}",ans);
}

ここでは、[i,j)という区間を伸ばしていくという発想にしています。

なぜこうするかというと、ループが始まった瞬間に[i,j)の合計はaccという不変式がほしいからです。これにより、初期値に一貫性を持てます。

そして、条件が成立するまで回すwhileループと、条件が成立していた場合の条件判定は全く独立させることにしています。こうすることで、jがNに突き刺さって動かない状態になった場合にも、他の部分に影響がなくなります。whileとifの間にフラグなどで依存関係があるとややこしくなります。