Pythonで連続する同じ値の要素をサブリストにまとめる方法

数値のリスト nums が与えられたとき、同じ値が連続して並んでいる要素をひとつのサブリストにまとめてパックするプログラムを作成します。ここで注意すべき点は、リスト内に1回しか出現しない要素であっても、それ単独のサブリストとして残す必要があるということです。

たとえば、入力が nums = [5, 5, 2, 7, 7, 7, 2, 2, 2, 2] の場合、出力は次のようになります。

[[5, 5], [2], [7, 7, 7], [2, 2, 2, 2]]

同じ値でも離れて登場する場合は別々のサブリストになるため、2 が2箇所に分かれている点に注目してください。

解き方の手順

• nums が空の場合は、空のリストを返す
• result を、先頭要素 nums[0] を含むリストで初期化する（つまり [[nums[0]]]）
• インデックス j を 0 に設定する
• i を 1 から nums のサイズまで繰り返す：
  • nums[i] が直前の要素 nums[i - 1] と異なる場合
    • result の末尾に新しい空のリストを追加する
    • j を 1 増やす
  • nums[i] を result[j] の末尾に追加する
• result を返す

実装例

それでは、上記の手順を実際のコードで確認してみましょう。

class Solution:
   def solve(self, nums):
      if not nums:
         return []
      result = [[nums[0]]]
      j = 0
      for i in range(1, len(nums)):
         if nums[i] != nums[i - 1]:
            result.append([])
            j += 1
         result[j].append(nums[i])
      return result

ob = Solution()
nums = [5, 5, 2, 7, 7, 7, 2, 2, 2, 2]
print(ob.solve(nums))

入力

[5, 5, 2, 7, 7, 7, 2, 2, 2, 2]

出力

[[5, 5], [2], [7, 7, 7], [2, 2, 2, 2]]

itertools.groupbyを使ったより簡潔な方法

Pythonの標準ライブラリ itertools.groupby を使うと、同じ処理をわずか数行で実現できます。groupbyは「隣接する同じキーの要素」をグループ化する関数なので、この問題にぴったりです。

from itertools import groupby

def solve(nums):
   return [list(group) for _, group in groupby(nums)]

nums = [5, 5, 2, 7, 7, 7, 2, 2, 2, 2]
print(solve(nums))
# 出力: [[5, 5], [2], [7, 7, 7], [2, 2, 2, 2]]

どちらの方法でも計算量は O(n) で、リストを一度走査するだけで済みます。自前でループを書きたい場合は最初の方法、可読性と簡潔さを重視する場合は groupby を使う方法がおすすめです。