Pythonでシャッフルされたキューを元の順序に復元するプログラム

問題の概要

ここに2次元のマトリックス（リスト）があります。各行は [height, count] という2つの値を持っており、height はその人の身長、count は「その人の前方にいる、身長が自分と同じかそれ以上の人の数」を表します。このキューがシャッフルされてしまったとき、元の並び順を復元するのが今回の課題です。

たとえば、入力が次のような場合を考えてみましょう。

このとき、期待される出力は次のとおりです。

結果を見ると、身長4と5の人は前方に自分以上の背の人がいないため先頭側に配置され、身長2の人の前には身長4と5の2人が立っているため、3番目に配置されていることがわかります。

解決のための手順

この問題は、次の手順で解くことができます。

1. N := マトリックスの行数を取得します。
2. マトリックスの行を「身長の昇順・カウントの降順」で並べ替えます。
3. ans := サイズNのリストを作成し、すべての要素を None（空）で初期化します。
4. マトリックスの各行から身長 h とカウント c を取り出し、以下の処理を行います。
   • temp := 0 とします。
   • ans の各インデックス i と値 num について、次を繰り返します。
     • temp >= c かつ num が None の場合：ans[i] := [h, c] を設定し、ループを抜けます。
     • num が None または num[0] >= h の場合：temp := temp + 1 とします。
5. 最後に ans を返します。

このアルゴリズムのポイントは、身長の昇順に並べ替えて背の低い人から順に確定させていく点です。すでに配置済みの人と、自分と同じ以上の身長の人だけをカウントすることで、各人が入るべき正しい位置を求められます。

実装例

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

class Solution:
    def solve(self, matrix):
        N = len(matrix)
        matrix.sort(key=lambda x: [x[0], -x[1]])
        ans = [None] * N

        for h, c in matrix:
            temp = 0
            for i, num in enumerate(ans):
                if temp >= c and num is None:
                    ans[i] = [h, c]
                    break

                if num is None or num[0] >= h:
                    temp += 1
        return ans

ob = Solution()
matrix = [
    [2, 2],
    [4, 0],
    [5, 0]
]
print(ob.solve(matrix))

入力

[[2, 2],[4, 0],[5, 0]]

出力

[[4, 0], [5, 0], [2, 2]]