Pythonで解く「最後の石の重さ」問題 ― アルゴリズムと実装をわかりやすく解説

問題の概要

それぞれ正の整数の重さを持つ石がいくつか与えられます。毎ターン、最も重い2つの石を選んで砕きます。2つの石の重さを x、y(x ≤ y)とすると、砕いた結果は次の2通りのいずれかになります。

• x = y の場合: 2つの石はともに完全に破壊されます。
• x ≠ y の場合: 重さ x の石は完全に破壊され、重さ y の石は新しい重さ y − x となります。

この操作を繰り返すと、最後には最大で1個の石が残ります。残った石の重さを求めてください(石が1個も残らない場合は 0 を返します)。

具体例

例として、石の重さが [2, 7, 4, 1, 8, 1] の場合を考えてみましょう。このとき答えは 1 になります。

1. まず最も重い 8 と 7 を砕くと、差は 1。配列は [2, 4, 1, 1, 1] になります。
2. 次に 4 と 2 を砕くと、差は 2。配列は [2, 1, 1, 1] になります。
3. 続いて 2 と 1 を砕くと、差は 1。配列は [1, 1, 1] になります。
4. 最後に重さ 1 の石同士を砕くと、両方とも消滅し、配列は [1] となります。

したがって、最後に残る石の重さは 1 です。

解法のアプローチ

この問題は、ルールどおりにシミュレーションを行うことで解けます。手順は以下のとおりです。

• 配列 W が空であれば 0 を返す。
• W の要素が1つだけであれば、その値(W[0])を返す。
• W の要素数が2以上である間、以下を繰り返す。
  • W をソートする。
  • s1 := W の末尾の要素(最大値)、s2 := 末尾から2番目の要素とする。
  • s1 = s2 であれば、両方の要素を W から取り除く。
  • そうでなければ、s1 := |s1 − s2| と計算し、末尾の要素を取り除いた後、新たな末尾の要素を s1 で置き換える。
• ループ終了後、W に要素が残っていればその値を、空であれば 0 を返す。

Pythonでの実装例

それでは、実際のコードを見てみましょう。

class Solution(object):
    def lastStoneWeight(self, stones):
        """
        :type stones: List[int]
        :rtype: int
        """
        if len(stones) == 0:
            return 0
        if len(stones) == 1:
            return stones[0]
        while len(stones) > 1:
            stones.sort()
            s1, s2 = stones[-1], stones[-2]
            if s1 == s2:
                stones.pop()
                stones.pop()
            else:
                s1 = abs(s1 - s2)
                stones.pop()
                stones[-1] = s1
        return stones[-1] if stones else 0

ob1 = Solution()
print(ob1.lastStoneWeight([2, 7, 4, 1, 8, 1]))

入力と出力

入力:

[2, 7, 4, 1, 8, 1]

出力:

1

計算量と効率化のポイント

上記の実装では、各ターンごとにソートを行っているため、時間計算量は O(n² log n)(n は石の個数)となります。より効率化したい場合は、優先度付きキュー(ヒープ)を利用するとよいでしょう。Python の heapq モジュールは最小ヒープしか提供していないため、重さを負の値にして格納することで最大ヒープとして扱えます。この方法なら各操作を O(log n) に抑えられ、全体の計算量を O(n log n) まで改善できます。

ヒープを使った実装例

import heapq

class Solution(object):
    def lastStoneWeight(self, stones):
        heap = [-s for s in stones]
        heapq.heapify(heap)
        while len(heap) > 1:
            y = -heapq.heappop(heap)
            x = -heapq.heappop(heap)
            if y != x:
                heapq.heappush(heap, -(y - x))
        return -heap[0] if heap else 0

状況に応じて、シンプルなソート方式とヒープ方式を使い分けることで、可読性とパフォーマンスのバランスを取ることができます。