Pythonで値kに基づいて連結リストのノードを並べ替えるプログラム

単方向連結リストとある値 k が与えられたとします。このとき、ノードを次の順序で並べ替える必要があります。まず値が k より小さいノードを先頭に集め、次に値が k と等しいノードを続け、最後にそれ以外（k より大きい）ノードを末尾に配置します。重要な制約として、ノード同士の相対的な順序は元のまま維持しなければなりません。

例えば、入力が L = [4, 3, 6, 6, 6, 10, 8]、k = 6 の場合、出力は [4, 3, 6, 6, 6, 10, 8] となります。4 と 3 は 6 未満なので先頭に、6 は 6 と等しいので中間に、10 と 8 は 6 より大きいため末尾に配置されますが、各グループ内の順序は元のリストのまま変わりません。

解決手順

この問題を解くには、以下の手順に従います。

• 値 0 を持つダミーノード less_head を作成し、ポインタ less を less_head に設定する
• 同様に、ダミーノード equal_head と greater_head を作成し、それぞれポインタ equal、greater を設定する
• cur を先頭ノードに設定する
• cur が null になるまで以下を繰り返す
  • cur の値が k より小さい場合：less の後ろに新しいノードを追加し、less を進める
  • cur の値が k より大きい場合：greater の後ろに新しいノードを追加し、greater を進める
  • それ以外の場合：equal の後ろに新しいノードを追加し、equal を進める
  • cur を次のノードに進める
• less の next を equal_head の next につなぎ、equal の next を greater_head の next につなぐことで、3つのリストを1つに連結する
• less_head の next を返す

ここでダミーノードを使う理由は、各グループの先頭への挿入処理を簡潔にし、null チェックを不要にするためです。また、このアルゴリズムはリストを一度だけ走査するため、計算量は O(n)、追加のメモリ使用量は O(n) で済みます。

理解を深めるために、以下の実装を見てみましょう。

実装例

class ListNode:
   def __init__(self, data, next=None):
      self.val = data
      self.next = next

def make_list(elements):
   head = ListNode(elements[0])
   for element in elements[1:]:
      ptr = head
      while ptr.next:
         ptr = ptr.next
      ptr.next = ListNode(element)
   return head

def print_list(head):
   ptr = head
   print('[', end="")
   while ptr:
      print(ptr.val, end=", ")
      ptr = ptr.next
   print(']')

class Solution:
   def solve(self, node, k):
      less_head = less = ListNode(0)
      equal_head = equal = ListNode(0)
      greater_head = greater = ListNode(0)
      cur = node
      while cur:
         if cur.val < k:
            less.next = ListNode(cur.val)
            less = less.next
         elif cur.val > k:
            greater.next = ListNode(cur.val)
            greater = greater.next
         else:
            equal.next = ListNode(cur.val)
            equal = equal.next
         cur = cur.next
      less.next = equal_head.next
      equal.next = greater_head.next
      return less_head.next

ob = Solution()
L = make_list([4, 3, 6, 6, 6, 10, 8])
k = 6
print_list(ob.solve(L, k))

入力

[4, 3, 6, 6, 6, 10, 8], 6

出力

[4, 3, 6, 6, 6, 10, 8]