コーディング面接: クイックセレクトを使用して重複を含む k 番目に大きい要素を見つける

問題と適用範囲

ソートされていない整数配列 nums と 1 ≤ k ≤ nums.length が与えられたとき、広義の降順（非増加順）で k 番目の要素を返します。重複は個別の位置を占めます。たとえば [5, 5, 4] において2番目に大きい値は、2番目に異なる値ではなく 5 です。配列の変更が許可されている場合、目標の平均時間計算量は O(n)、追加の空間計算量は O(1) です。コーディングを始める前にその前提を述べてください。

面接官がテストしていること

優れた回答では、「k 番目に大きい」を昇順のインデックス target = n-k にマッピングし、パーティションはピボットの周囲に境界を確立するだけでよく、もう一方の側をソートする必要はまったくないことを説明します。また、重複、k=1、k=n、ソート済み入力の処理、およびランダム化された平均計算量と最悪ケースの保証の違いについても扱います。

コーディング前の確認事項

1. 入力を変更してもよいですか？ インプレースのパーティションは追加の空間計算量が O(1) ですが、入力を保持するには O(n) のコピーが必要です。
2. これは k 番目の位置ですか、それとも k 番目に異なる値ですか？ 通常は位置を数える問題文ですが、異なる値の選択には別の重複処理が必要です。
3. データはストリームとして届きますか？ クイックセレクトは1つの実体化された配列を対象とします。サイズ k の最小ヒープを使用すると、ストリームに対して O(n log k) の処理が可能になります。
4. 決定論的な最悪ケースの限界は必須ですか？ ランダム化クイックセレクトは平均 O(n) です。厳密な最悪ケースを主張するには、中央値の中央値（median-of-medians）または標準ライブラリの保証が必要です。

推奨される解法と導出

ピボットより小さい値、等しい値、大きい値の3つに分割する3-wayパーティショニングを使用します。変換された昇順インデックス target について、ターゲットが lt の左側にある場合は左の間隔で続行し、gt の右側にある場合は右の間隔で続行し、[lt, gt] 内にある場合はピボットを返します。等しい要素の範囲を一度に扱うことで、すべての要素が等しい入力に対して、1つずつ要素を繰り返し破棄する代わりに1回のスキャンで処理を完了できます。

import random

def kth_largest(nums: list[int], k: int) -> int:
    if not 1 <= k <= len(nums):
        raise ValueError("k out of range")
    target = len(nums) - k
    left, right = 0, len(nums) - 1
    while left <= right:
        pivot = nums[random.randint(left, right)]
        lt, i, gt = left, left, right
        while i <= gt:
            if nums[i] < pivot:
                nums[lt], nums[i] = nums[i], nums[lt]
                lt += 1; i += 1
            elif nums[i] > pivot:
                nums[i], nums[gt] = nums[gt], nums[i]
                gt -= 1
            else:
                i += 1
        if target < lt:
            right = lt - 1
        elif target > gt:
            left = gt + 1
        else:
            return pivot
    raise RuntimeError("unreachable")

各イテレーションでは、現在の間隔を1回スキャンします。ピボットによって間隔が一定の割合で縮小される場合、T(n)=T(cn)+O(n) により平均 O(n) が得られます。極端な値を繰り返し選択すると、最悪のケースでは依然として O(n²) になります。反復処理の形式にすることで再帰の深さを回避し、O(1) の追加空間を使用します。

代替案とトレードオフ

完全なソートは検証が最も簡単で、計算量は O(n log n) であり、配列が小さい場合や後で完全な順序が必要な場合に合理的です。サイズ k の最小ヒープは入力を保持し、時間計算量は O(n log k)、空間計算量は O(k) となり、ストリームや k が n よりもはるかに小さい場合に適しています。C++ の std::nth_element は、平均線形時間計算量で同じパーティションセマンティクスを公開します。選択された位置のどちらの側もソートしません。

失敗モード、境界条件、および反例

• target = k-1 と記述すると、k 番目に小さい値が検出され、要求された順序と逆になります。
• 等しい要素を1つしか破棄しない2-wayパーティションは、[7, 7, 7, ...] に対して O(n²) かかる場合があります。3-wayパーティションは等しい要素の範囲を一度に消費します。
• 常に末尾の要素を選択すると、ソート済みおよび逆順ソート済みの入力でパフォーマンスが低下する可能性があります。ランダム化によってその可能性は低くなりますが、漸近的な最悪ケースの上限は変わりません。
• 「k 番目に異なる最大値」の場合、等しい要素の範囲をカウントまたは削除せずにこの停止条件を再利用することはできません。
• 空の配列、k=0、または k>n は、インデックスエラーによって無効な問題を隠してしまうのではなく、境界で拒絶してください。

テストと検証のチェックリスト

ランダム化されたケースを sorted(nums)[-k] と比較します。すべての要素が等しい値、負の数と重複、k=1、k=n、ソート済み入力、および逆順ソート済み入力を含めます。配列の変更が許可されている場合は、配列全体の順序ではなく結果をアサートします。再現性を確保するために乱数シードを固定し、n の増加に伴う比較回数を記録します。たまたま1回うまくいっただけでは計算量の証明にはなりません。

フォローアップの質問

最悪ケースで線形時間を保証するにはどうすればよいですか？

中央値の中央値（median-of-medians）ピボットを選択することで、毎ラウンド一定の割合が削除され、最悪ケースの時間計算量が O(n) になります。これにはより大きな定数係数が伴うため、本番コードでは通常、ランダム化された選択や標準ライブラリの実装が選ばれます。

k 番目に小さい要素に変更するにはどうすればよいですか？

昇順のパーティションを維持したまま target = k-1 を使用します。最大値を求める定式化を target=n-k として保持する方が、配列を反転するよりもわかりやすい場合が多くあります。

挿入と多数の順位クエリをサポートするにはどうすればよいですか？

単発のクイックセレクトでは、クエリごとに再スキャンが行われます。固定された1つの k に対しては、サイズ k の最小ヒープを維持します。任意の順位クエリに対しては、サブツリーのサイズで拡張された平衡木を検討し、更新とクエリの比率に基づいて選択します。