[Python] 競プロをきっかけにアルゴリズムを勉強し始める - スタックとキュー [Paiza]

この記事は第2回目です。こちらの記事に、私がこの記事を書くことになった経緯がありますので、よければそちらのイントロからお読み下さい。
内容は Python 100% です。

スタックとキューとは

今回は、まず基本のデータ構造と言える スタック stack と キュー queue について解説しようと思います。お恥ずかしながら、私もこうやってアルゴリズムの勉強を真面目に始めるまでは「知っているけど使ったことないな」というレベルでした。私の専門である自然言語処理でも、文章のパーシング（構文解析）などではスタックを使うらしいのですが、自分で実装はしたことがありません。

説明に関しては、私が書くよりもこちらの記事を読んで頂いた方がわかりやすいかと思います（他力本願）。非常に良くまとまっていて勉強になりますので、まず読んで下さい。この中に載っている、「ヒストグラム中の最大面積」に似た問題を解いた時、私はまんまと $O(n^2)$ のコードを書いて制限時間に引っかかったことがあります。

スタックとキューはどちらも配列です。違いは、後入れ先出し (LIFO: Last In First Out) か、先入れ先出し (FIFO: Fisrt In First Out) か、という点です。

「でもそれって普通のリストを使って、先頭か最後を操作すればいいだけじゃないの？」と思った方、はいそうです、私もそう思っていた時期がありました。

collections.deque

Python の組み込みモジュールである collections の中に、deque というデータ型があり、これの説明を引用します。

Deque とは、スタックとキューを一般化したものです (この名前は「デック」と発音され、これは「double-ended queue」の省略形です)。Deque はどちらの側からも append と pop が可能で、スレッドセーフでメモリ効率がよく、どちらの方向からもおよそ O(1) のパフォーマンスで実行できます。

list オブジェクトでも同様の操作を実現できますが、これは高速な固定長の操作に特化されており、内部のデータ表現形式のサイズと位置を両方変えるような pop(0) や insert(0, v) などの操作ではメモリ移動のために O(n) のコストを必要とします。

要するに、普通のリストに .insert() したり .pop() したりする時に、例えそれが先頭であったとしても、リスト全体を作り替える羽目になるため、$O(n)$ の時間がかかるよ ということです。

一方で、スタックでの LIFO やキューでの FIFO といった操作をするのに特化した配列型が、collections.deque です。こちらは、先頭や最後へのデータの出し入れは $O(1)$ 時間でできます。collections の中では Counter や OrderedDict なんかは多用していましたが、deque は何年も Python 触ってて初めて使いました、ごめんなさい。

ということで、実際にどう速いのか notebook 上で試してみましょう。from collections import deque した上で、以下の2つを実行してみて下さい。

リスト

%%time
L = []
for _ in range(20_000_000):
    L.append(i)

CPU times: user 997 ms

deque

%%time
L = deque()
for _ in range(20_000_000):
    L.append(i)

CPU times: user 1.03 s

あれれ、ほぼ誤差レベルでしか変わりませんね...じゃあやっぱりリストでいいじゃん！となりそうですが、今度はこちらでお試し下さい。どちらも先頭から 1000 回要素を取り除いています。

リスト

%%time
L = list(range(20_000_000))
for _ in range(1000):
    L.pop(0)

CPU times: user 5.45 s

deque

%%time
L = deque(range(20_000_000))
for _ in range(1000):
    L.popleft()

CPU times: user 234 ms

右側への .append() や pop(-1) だけではほとんど差が出ませんが、このように先頭から取り除くという場合においては、配列が大きいと圧倒的な差が出ます。つまり、後入れ先出しのスタックならリストでも問題になりにくいが、先入れ先出しのキューならリストはよろしくないということです。

皆さん以下のようなコードを書いてしまった経験はありませんか？私は多分何百回とあるはずです。

for x in ...:
    L.append(x)
    L = L[1:]

このようなコードがある場合は、deque で置き換えた方がパフォーマンスが向上します。もちろん、配列が10000個程度のものなら誤差レベルなので普通のリストで十分ですが、後々出てくる幅優先探索なんかではこれを使わないと死ぬ場面も出てきます。また、先頭・最後以外へのアクセスが必要な場合は、リストでも deque でも大差ありません。ちなみに、こちらの記事に各操作の時間計算量が詳しく載っていました。

queue.Queue

実は、Python にはもう一つの組み込みモジュールである queue というものが存在します。

こちらはその名の通り、キューそのものです。具体的な違いはスレッドセーフが保証されているかという点ですが、競プロならマルチスレッドで処理することなどもないですし、とりあえず deque を使っておけば速いし間違いない です。ということで今後も基本は deque のみ使っていきます。

スタックを使ったアルゴリズム

deque を使うと先頭部分での出し入れで時間計算量を節約できることは理解して頂けたかと思いますが、それは本質ではありません。スタックとキューとはあくまでデータ構造であり、リストを使って実装するか deque を使って実装するかは時間さえ気にしなければどちらでもいいのです。「どのような場面でスタックやキューを使う必要があるか？」という考え方こそが重要かと思います。

ではまずスタックの例から挙げます。一番有名なのは、おそらく逆ポーランド記法かと思います。

• 演算される2つの数字を先に置き、その後ろに演算子を置く
• 後に計算する演算子ほど、後方に置く

というものです。

この Wiki にも載っている、3 4 + 1 2 - * という例を使います。これは (3+4)*(1-2) と同値です。これを段階を追って計算すると以下のようになります。

1. + の前の2つの数をとり、3 4 + を計算することで、7 1 2 - * となる
2. - の前の2つの数をとり、1 2 - を計算することで、7 -1 * となる
3. * の前の2つの数をとり、7 -1 * を計算することで、-7 となる

演算子が来たときに、その2つ前をとるというのは、後方から取得するスタックであると言えます。これを Python を使って順番に推移を表すと以下のように表せます。

equation = [3, 4, '+', 1, 2, '-', '*']
stack = []

for x in equation:
    ...

x = 3

equation = [4, '+', 1, 2, '-', '*']
stack = [3]

x = 4

equation = ['+', 1, 2, '-', '*']
stack = [3, 4]

x = '+'

equation = [1, 2, '-', '*']
stack = [7]  # 3 + 4  

x = 1

equation = [2, '-', '*']
stack = [7, 1]

i = 2

equation = ['-', '*']
stack = [7, 1, 2]

x = '-'

equation = ['*']
stack = [7, -1]  # 1 - 2

x = '*'

equation = []
stack = [-7]  # 7 * -1

このように、式から一つずつ取り出し、それをスタックに入れて、演算子が来た場合は後ろから2つ取り出して計算し、またスタックに入れ直すという作業を繰り返しています。ちなみに元の式の要素が減っているのは見た目上わかりやすくしただけのことで、実際は for ループするだけなので操作する必要はありません。

Paiza の問題例

こちらに、まさに逆ポーランド記法の練習問題があります。

私の解答例を載せておきますので参考にして下さい。今回はスタックなので deque ではなくリストを使って実装しています。

N = int(input())
equation = input().split() # 例: 1 2 + 3 4 + -

stack = []

for x in equation:
    if x.isdigit():
        stack.append(int(x))
    else:
        n2 = stack.pop(-1)
        n1 = stack.pop(-1)
        
        if x == '+':
            stack.append(n1 + n2)
        elif x == '-':
            stack.append(n1 - n2)
            
print(stack[0])

ちなみに、文字列で書かれた式を評価する eval() という関数を使って、以下のように書くことも可能です。この場合、演算子が増えても if - elif 分岐なしで統一的に処理できます。

N = int(input())
equation = input().split() # 例: 1 2 + 3 4 + -

stack = []

for x in equation:
    if x.isdigit():
        stack.append(x)
    else:
        n2 = stack.pop(-1)
        n1 = stack.pop(-1)
        stack.append(eval(f'{n1}{x}{n2}'))
        
print(stack[0])

もう一問紹介します。こちらは括弧列と言われるものです。これは木構造とも関連しますね。

私の VSCode ではこのように括弧列に色が付きますが、これ同様に括弧が正しいか判定しろという問題です。S問題でもこれの発展系的な問題がありました。

これは逆ポーランド記法同様に左括弧をどんどんスタックに入れていき、右括弧がきたら一番最後の左括弧を取り出せばいいだけです。正しくない括弧列になる場合は、

• 右括弧 ) が来たときに、スタックが空である
• 全部終わった時にスタックに残っている

の2パターンです。実装法はたくさんありますが、以下の例は欲張って try - except と for - else を使ってみました。for - else は、for ループが最後まで終わった時に行われる処理なので、break した場合にはスキップされます。

N = int(input())
S = input() # 例 ((((())())()))

stack = []

for c in S:
    if c == ')': # 右括弧が来たとき、スタックから取り出す
        try:
            stack.pop(-1)
        except:
            print('No') # スタックが空だとエラーなので、break して終了
            break
    elif c == '(': # 左括弧が来たとき、スタックに追加
        stack.append(c)
else:
    if len(stack) > 0: # 全て終わった時、スタックに残っている
        print('No')
    else:
        print('Yes')

キューを使ったアルゴリズム

次にキューです。定番の幅優先探索問題などはそのうちやるとして、Paiza のこの問題を見てみましょう。

要するに「連続する X 個の要素の和の最大」を出すだけなので、実装自体は全然難しくありません。ですが、何やら問題文の中にも、「タイムオーバーになってしまうことがあります」という警告があります。この解答例をもう少し短く書くと以下のような感じです。

二重 for ループ

n, x = map(int, input().split())
nums = list(map(int, input().split()))

left_num = -1
max_sum = 0

for i in range(n - x + 1):    # 区間左端のインデックス範囲
    sum_x = sum(nums[i:i+x])  # X個の和を取る
    if sum_x > max_sum:       # 最大値の更新 
        left_num = nums[i]
        max_sum = sum_x

print(max_sum, left_num)

非常にシンプルですね。sum() は実質 for ループをしているのと同じことなので、「二重 for ループ」と表記しました。じゃあ実行してみましょう。

案の定タイムオーバーです。問題の警告文にも書いてありますが、sum() はその都度 X 個の足し算をしているので、外側の for の数 N と合わせて N * X 回の演算が必要になります（$O(NX)$ の計算量）。テスト3ではそれぞれ N=5000000, X=10443 なので、52,215,000,000 回となります。コンピュータの1秒あたりの演算回数は、処理にもよりますがせいぜい $10^8$ 程度が限界です。どう考えても間に合うわけがありません。

そこでこのような方法を考えます。

• 先頭から X 個の数を配列(キュー)に保存
• その X 個の和を計算し、sum_x として記録しておく
• X+1 番目の数を配列に入れるとともに、先頭から一個取り出し、その差分で sum_x を再計算
• sum_x が最大を更新するかどうかを判定

これであれば、すでに計算してある直前の X 個の和に対して、「新しく追加した1個をプラス」「取り出した1個をマイナス」という2回の演算だけで済むことになり、X の大きさに関わらず 2N 回となります。deque を使った実装だと以下のような感じです。

キューを用いた方法

from collections import deque

n, x = map(int, input().split())
nums = list(map(int, input().split()))

# キューの作成と、記録用変数の宣言
# 和の最大値と左端の値は、初期値を設定
queue = deque(nums[:x])
sum_x = sum(queue)
max_sum = sum_x
left_num = nums[0]

for right in nums[x:]:
    left = queue.popleft()        # 左から1つ取り出す
    queue.append(right)           # 右に1つ追加
    sum_x = sum_x - left + right  # 差分から和を計算
    if sum_x > max_sum:           # 最大値の更新
        max_sum = sum_x
        left_num = queue[0]

print(max_sum, left_num)

この方法だと、タイムオーバーしたテスト3も0.27秒で終了します。ちなみに、 deque を使ってキューを作らなくても、単に配列の特定の位置にアクセスすればいいだけなので、以下のようにも書けます。

記録用変数のみの方法

n, x = map(int, input().split())
nums = list(map(int, input().split()))

sum_x = sum(nums[:x])
max_sum = sum_x
left_num = nums[0]

for left_i in range(n - x):
    sum_x = sum_x - nums[left_i] + nums[left_i+x]
    if sum_x > max_sum:
        max_sum = sum_x
        left_num = nums[left_i+1]

print(max_sum, left_num)

こちらの方が時間も若干短縮されるだけでなく、空間計算量(メモリ消費)も少なくなります。とは言ってもそこまで影響が出るレベルではないでしょうが。ただ、range の範囲を決める時に、私は n-x? n-x+1? n-x-1? のようにしばらく迷ってしまいましたね。その点上記のキューを用いる方法では、「最初に num[:x] の和を計算しておく」「num[x:] から一つずつ取り出す」と非常に理解しやすく、個人的には好みです。

さらに累積和というものを使って解く方法もありますが、それは今回の場合は少し冗長なので、また次の機会に話そうと思います。

エスカレーター

さらにもう一問、エスカレーターに関する問題を紹介します。片側から追い抜いて行かない限りは典型的なキューの構造です。

以下のような考え方でできそうです。

• エスカレーターを表すキューを用意
• 人が乗っている場所は 1、いない場所は 0 で表す
• 現在の時刻から次に社員が乗ってくる時刻まで、先頭からの取り出し (dequeue) と最後尾への追加 (enqueue) を繰り返し行う。

このやり方の場合、t = 1 の時刻に社員が乗ってきて、そこからしばらく空いて t = 1001 まで誰も乗ってこない場合、何百回と無駄な取り出し＆追加をすることになるので効率が悪いと思うかもしれません。ただ deque を使った場合は $O(1)$ の定数時間でできますし、「人が頻繁に乗り降りする場合」「人の乗り降りが少ない場合」のどちらであっても、同じコードで同じように動くというのは大事なことです。

ここで、deque のキーワード変数である maxlen というものを活用します。これはキューの最大長を決めるもので、もし最大長を超えて追加された場合には、自動的に先頭の要素から削除されます。つまり、.append() だけで全てが行えるということです。今回のように、取り出された値自体に興味がない場合は非常に有効ですね。

from collections import deque

N, K = map(int, input().split()) # 社員数、エレベーター長
times = list(map(int, input().split())) # 社員が乗ってくる時刻

L = deque([0]*K, maxlen=K)  # K の固定長のキュー
now = 0  # 現在時刻

for t in times:
    for _ in range(t-now-1): # 現在時刻から、乗ってくる時刻の1秒前まで
        L.append(0)  # 0 を右に追加、自動的に左から1つ削除
    L.append(1)   # 社員が乗ってくるので 1 追加
    print(sum(L)) # エレベーター上の人数を計算
    now = t  # 現在時刻の更新

for が二重になっていますが、結局は時刻通りに辿っているだけのことなので時間計算量としては、最後の時刻に依存するということになります。ちなみに、内側の for を使わずに .extend() というリストと同じメソッドを使って以下のようにスマートに書くこともできます。

from collections import deque

N, K = map(int, input().split()) # 社員数、エレベーター長
times = list(map(int, input().split())) # 社員が乗ってくる時刻

L = deque([0]*K, maxlen=K)  # K の固定長のキュー
now = 0  # 現在時刻

for t in times:
    L.extend([0]*(t-now-1)) # 現在時刻から、乗ってくる時刻の1秒前まで 0 追加
    L.append(1)   # 社員が乗ってくるので 1 追加
    print(sum(L)) # エレベーター上の人数を計算
    now = t  # 現在時刻の更新

スタック・キューの練習問題

上記のスタック・キューの問題はこちらにまとまっています。

• 基本編最終問題「箱とボール」
• 応用編「積読」
• 応用編「満員エスカレーター」

このうち、以上の3つをピックアップしました。私の解答例を載せておきますので参考にしてください。

箱とボール

N = int(input())
balls = map(int, input().split())

box = []  # スタックの作成

for ball in balls:
    box.append(ball)
    while len(box) >=2 and box[-1] == box[-2]: # 最後の2つが同じなら
        box.pop(-1)  # 1つを取り出し
        box[-1] *= 2  # 残った1つを2倍   
        
print(*box[::-1], sep='\n')  # 逆順で取り出す

積読

n = int(input())
actions = [input().split() for i in range(n)]

stack = []

for action, page in actions:
    page = int(page)
    
    if action == 'buy_book':
        stack.append(page)
    else:
        while page > 0:
            if stack[-1] > page:
                stack[-1] -= page
                page = 0
            else:
                page -= stack.pop(-1)
                
print(len(stack))
print(sum(stack))

満員エスカレーター

from collections import deque

n, limit = map(int, input().split())
rides = [input().split() for i in range(n)]

esc = deque()
total_weight = 0

for ride in rides:
    action = ride[0]  # ride or getoff
    num = int(ride[1])
    
    if action == 'ride':
        weights = [int(x) for x in ride[2:]]
        for w in weights:
            if total_weight + w <= limit:
                esc.append(w)
                total_weight += w
    else:
        for i in range(num):
            total_weight -= esc.popleft()
            
print(len(esc))
print(total_weight)

英語ならば、LeetCode がジャンル別・難易度別にまとまっているので参考になります。

この Stack の最初の問題、20. Valid Parentheses なんかはまさに上述の括弧列と同じような問題です。私の解答例です。

class Solution:
    def isValid(self, string: str) -> bool:
        stack = []
        for c in string:
            if c in ')}]':
                pair = {')':'(', ']':'[', '}':'{'}[c]
                if len(stack) == 0 or stack[-1] != pair:
                    return False
                else:
                    stack.pop(-1)
            else:
                stack.append(c)

        return len(stack) == 0

次回は再帰について話そうと思います。