LeetCode 232: Implement Queue using Stacks - 今日も窓辺でプログラム

LeetCodeに、スタックを使ってキューを実装せよという問題がありました。
Implement Queue using Stacks - LeetCode Articles

pushが計算量 $O(1)$ 、popも償却計算量が $O(1)$ となる方法が紹介されていたので、まとめておきます。
償却計算量が何かについては、私の以前の記事を参照してください。
償却計算量とは何か - 今日も窓辺でプログラム

LeetCodeというサイト自体について紹介している記事もあるので、知らない方はチェックしてみてください。
LeetCode: コーディング面接に向けた練習に使えるサイトの紹介 - 今日も窓辺でプログラム

最初に思いついた方法

最初に思いついて実装したコードは、これでした。

class Queue(object):
    def __init__(self):
        """
        initialize your data structure here.
        """
        self.s1 = []
        self.s2 = []


    def push(self, x):
        """
        :type x: int
        :rtype: nothing
        """
        self.s1.append(x)


    def pop(self):
        """
        :rtype: nothing
        """
        if len(self.s1) == 0:
            return

        while len(self.s1) > 1:
            self.s2.append(self.s1.pop())
        self.s1.pop()

        while len(self.s2) > 0:
            self.s1.append(self.s2.pop())

    def peek(self):
        """
        :rtype: int
        """
        if len(self.s1) == 0:
            return

        while len(self.s1) > 1:
            self.s2.append(self.s1.pop())
        val = self.s1.pop()
        self.s2.append(val)

        while len(self.s2) > 0:
            self.s1.append(self.s2.pop())

        return val


    def empty(self):
        """
        :rtype: bool
        """
        return len(self.s1) == 0

2つのスタックを用意して、popやpeekをするときに既にスタックに積まれている要素をもう片方のスタッフにどけておき、スタックの底の要素を取り出して、また元に戻す、という動作です。

これでも正しくキューの動きをするのですが、popやpeekで毎回 $O(n)$ の時間がかかるのが難点です。

償却計算量O(1)の方法

アルゴリズム

2つのスタックを使って、pushが $O(1)$ 、popも償却計算量が $O(1)$ となる方法があるようです。
一言でいうと、「洗い終わったタオルを古いものから順番に使う」アルゴリズムです。

洗濯が終わったタオルを順に積み上げていくと、一つの山ができます。これをスタック１と呼びます。単純にスタック１の上からとってしまうと新しいタオルから使うことになってしまうので、タオルを使いたいときにはスタック１をひっくり返して横にずらして、もう一つのスタック２を用意します。

タオルを使うときはスタック２から使い、新しく選択したタオルはスタック１に置いていく。こうすると、スタック2つでキューの動作を実現できます。実装としてはこんな感じになります。

class Queue(object):
    def __init__(self):
        """
        initialize your data structure here.
        """
        self.s1 = []
        self.s2 = []
        self.front = None

    def push(self, x):
        """
        :type x: int
        :rtype: nothing
        """
        if len(self.s1) == 0:
            self.front = x
        self.s1.append(x)

    def pop(self):
        """
        :rtype: nothing
        """
        if len(self.s2) > 0:
            self.s2.pop()
        else:
            # Move s1 to s2
            while len(self.s1) > 0:
                self.s2.append(self.s1.pop())
            self.s2.pop()

    def peek(self):
        """
        :rtype: int
        """
        if len(self.s2) > 0:
            return self.s2[-1] # return top of stack
        elif len(self.s1) > 0:
            return self.front
        return None

    def empty(self):
        """
        :rtype: bool
        """
        return len(self.s1) == 0 and len(self.s2) == 0

popの計算量の評価

popのでは、スタック2がからでない場合はスタック2の要素をpopしているだけなので $O(1)$ です。
一方で、スタック2が空の場合は、スタック1の要素をすべてスタック2に移す作業（タオルをひっくり返す作業）が発生するので、計算量は $O(n)$ となります。なので最悪計算量は $O(1)$ となるのですが、償却計算量はどうなるのでしょうか。

償却計算量では、n回の連続した操作を考えます。キューなので、まずpushをn回、そしてpopをn回行った場合の計算量を見てみます。
pushは $O(1)$ だったので、n回行った場合は $O(n)$ です。
popは最初の1回はタオルをひっくり返す作業が発生するので $O(n)$ となりますが、残りの $n - 1$ 回に関しては $O(1)$ となります。

つまり、pushをn回、popをn回を連続で行った場合の計算量は
${\displaystyle O(1) * n + O(n) + O(1) * (n - 1) = O(n) }$
となります。これを合計操作回数 $2n$ で割った値が償却計算量となるので、償却計算量は $O(1)$ となります。

参考記事

Implement Queue using Stacks - LeetCode Articles