Pythonで株価が利益になるまでの最短待機日数を求めるプログラム

問題概要

ある企業の日々の株価が、時系列順に並んだリストとして与えられているとします。ここで、元のリストと同じ長さの新しいリストを作成することを考えます。新しいリストのインデックス i の値は、「その日の株価を上回る価格が現れるまでに何日待てばよいか」という最小日数を表します。もし将来どの日を見ても利益を得られる見込みがない場合は、その値は 0 とします。

例えば、入力が prices = [4, 3, 5, 9, 7, 6] の場合、出力は [2, 1, 1, 0, 0, 0] となります。

• インデックス 0(価格 4):2 日後に価格 5 となり利益を得られます
• インデックス 1(価格 3):1 日後に価格 5 となり利益を得られます
• インデックス 2(価格 5):1 日後に価格 9 となり利益を得られます
• インデックス 3 以降:それ以降にこれらの価格を上回る日は存在しないため 0 です

解法のアプローチ:単調スタック

この問題は「単調スタック(Monotonic Stack)」というテクニックを使うことで効率的に解くことができます。各日について未来を総当たりで調べると O(n²) かかりますが、スタックを活用すれば O(n) で処理できます。

具体的な手順は以下の通りです。

1. ans := prices と同じサイズのリストを用意し、すべて 0 で初期化する
2. q := 空のリスト(スタック)を用意する。スタックには (インデックス, 価格) のペアを格納する
3. prices の各インデックス i と価格 p について、以下を繰り返す:
   • q が空でなく、かつ p が q の末尾要素の価格より大きい間、次を繰り返す:
     • j := q の末尾要素のインデックス
     • ans[j] := i - j(j 日目から i 日目までの待ち日数を記録)
     • q から末尾要素を削除する
   • (i, p) を q の末尾に追加する
4. 最後に ans を返す

このアルゴリズムでは、現在の価格がスタック内の過去の価格を上回るたびに、それらの「待ち日数」が確定していきます。各要素は最大 1 回しかプッシュ・ポップされないため、全体の計算量は O(n) に抑えられます。

Pythonでの実装例

class Solution:
    def solve(self, prices):
        ans = [0 for _ in prices]
        q = []
        for i, p in enumerate(prices):
            while q and p > q[-1][1]:
                j = q[-1][0]
                ans[j] = i - j
                q.pop()
            q.append((i, p))
        return ans

ob = Solution()
prices = [4, 3, 5, 9, 7, 6]
print(ob.solve(prices))

入力

[4, 3, 5, 9, 7, 6]

出力

[2, 1, 1, 0, 0, 0]

まとめ

本記事では、単調スタックを用いて「利益が出るまでの最短待機日数」を効率的に求める方法を紹介しました。全探索では O(n²) となる問題を、O(n) の時間計算量で解ける点がこの手法の大きな魅力です。株価の分析だけでなく、「次により大きい要素(Nearest Greater Element)」系の問題全般に応用できる汎用的なテクニックなので、ぜひ覚えておきましょう。