キューブを積み上げられるかどうかを判定するPythonプログラム（deque＋貪欲法）

問題の概要

配列 nums には、n 個の異なるサイズのキューブの一辺の長さが横一列に並んで格納されています。これらのキューブをすべて使って、縦に一本の塔を積み上げることを考えます。

積み上げの際には、次のルールを守る必要があります。

• i 番目のキューブが j 番目のキューブの上に載るとき、下側（j 番目）のキューブの一辺の長さは、上側（i 番目）のキューブの一辺の長さ以上でなければならない。

さらに重要な制約として、キューブを取り出せるのは列の左端か右端のみで、中央から直接取ることはできません。この条件のもとで、すべてのキューブを積み上げられるかどうかを判定するのが本記事のテーマです。

具体例

たとえば、入力が nums = [1,2,3,7,8] の場合、出力は True になります。右端から順に「8 → 7 → 3 → 2 → 1」の順でキューブを取っていけば、大きいものほど下に配置されるため、ルールどおりに積み上げることができるからです。

一方、[4,5,2] のような配列ではどうでしょうか。最初に左端の 4 を取ると、残りの両端は 5 と 2 となり、5 は 4 より大きいため積めません。逆に右端の 2 を先に取っても、残った 4 と 5 はどちらも 2 より大きいため失敗します。この場合の出力は False です。

アルゴリズムの考え方

この問題は、両端キュー（deque）を使った貪欲法（グリーディ法）で効率よく解けます。ポイントは次のとおりです。

• 毎回、残っている列の左端と右端を比較し、大きい方（同値なら左端）を次に積むキューブとして選ぶ。
• 直前に積んだキューブより大きいキューブしか残っていない場合は、それ以上積めないため失敗（False）と判定する。

なぜ「大きい方を優先して取る」のが正しいのかというと、次に積めるキューブは必ず直前のキューブ以下でなければならないため、大きいキューブを後回しにすると、最終的に積めずに残ってしまう可能性があるからです。したがって、常に大きい方を先に消費する戦略が最適となります。

処理の手順

1. n を nums の要素数とする。
2. nums の要素から両端キュー d を作成する。
3. flag を True、prev を 0 で初期化する。
4. d が空になるまで、次の処理を繰り返す。
   • first を d の先頭要素、last を d の末尾要素とする。
   • prev が 0 以外で、first または last が prev より大きい場合は、flag を False にしてループを抜ける。
   • first が last 以上であれば、d の先頭要素を取り出して prev に代入する。
   • そうでなければ、d の末尾要素を取り出して prev に代入する。
5. flag が True なら True を、そうでなければ False を返す。

Python実装例

from collections import deque

def solve(nums):
    n = len(nums)
    d = deque(nums)
    flag = True
    prev = 0
    while d:
        first = d[0]
        last = d[-1]
        if prev != 0 and (first > prev or last > prev):
            flag = False
            break
        if first >= last:
            prev = d.popleft()
        else:
            prev = d.pop()
    if flag:
        return True
    else:
        return False

nums = [1,2,3,7,8]
print(solve(nums))

実行結果

入力：

[1,2,3,7,8]

出力：

True

計算量

各ループでキューブを1つずつ取り除いていくため、時間計算量は O(n) です。また、元の配列のコピーとして deque を保持するため、空間計算量も O(n) となります。要素数が増えても線形時間で判定できる、非常に効率的なアプローチです。