Sort Colors

ライブラリのソートを使えば O(n log n) で済みますが、問題の条件でそれは禁止されています。面接官は、値が3種類しかないという事実を踏まえてどうするかを見たいからです。そこで基本となるのは、自分で書くソートです。正しく実装するのが最も簡単なのはバブルソートです。配列を順に見ていき、隣り合う2つの値の順序が逆なら入れ替えます。

1回の走査で、そのときに見つけた最大値が泡のように上昇して最後まで移動します。1回目の走査後には最後の位置が確定し、2回目の走査後には最後の2つの位置が確定するため、n-1 回走査すれば配列全体が整列します。[2, 1, 0]の場合、1回目の走査で2が最後まで移動し、[1, 0, 2]になります。2回目の走査では1と0が入れ替わります。

遅いのは、各走査でまだ位置が確定していないすべてのペアを比較するからです。合計の比較回数はおよそ n²/2 回です。n = 1.5 × 10^4 では比較回数が 10^8 回を超え、さらに最初に順序が逆になっているペアごとに入れ替えも発生します。その作業のどれも、値が3種類しかないという事実を活用していません。

バブルソートは隣り合う要素の比較にすべての時間を費やしますが、どの値があるかはすでにわかっています。配列に0が2個、1が3個、2が2個あるなら、何も動かす前から答えは決まっています。0が2個、1が3個、2が2個です。重要なのは個数だけです。

そこで、配列を一度読み取り、それぞれの値を数えます。次に、先頭から配列を上書きします。count[0]個の0、続いてcount[1]個の1、そしてcount[2]個の2です。これが計数ソートです。等しい値は互いに入れ替えても同じなので、ここでは安全に使えます。1はどれも1なので、元の順序を保つ必要はありません。

これで2回の走査と3つのカウンターを使い、時間計算量はO(n)、空間計算量はO(1)です。条件を満たしており、色が多い場合には自然な解法です。この問題でよく知られている追加の問いは、配列を一度だけ読み取りながら同じことができるかどうかです。

読み進めながら3つの領域を作ります。先頭に0、続いて1、末尾に2を置き、1と2の間に未読部分を残します。3つのインデックスで境界を示します。lowより前はすべて0、lowからmidの直前まではすべて1、highより後ろはすべて2で、nums[mid]からnums[high]まではまだ未読です。

nums[mid]を読みます。1はすでに自分の領域にあるので、midを進めます。0は先頭に属するので、nums[low]と交換し、lowとmidの両方を進めます。lowから戻ってくる値は1（または、まだ1が見つかっていない場合は同じ0）なので、すでに正しい位置にあります。2は末尾に属するので、nums[high]と交換してhighを戻しますが、highから来た値はまだ読んでいないため、midはそのままにします。

各ステップでmidを進めるかhighを戻すため、未読部分は毎回1セルずつ小さくなり、ループはnステップ後に終了します。[2, 0, 2]を追ってみましょう。最初の2を最後の2と交換し、highは1になります。インデックス0にはまだ2があり、これを0と交換するとhighは0になります。インデックス0には今度は0があり、これはそのままなので、結果は[0, 2, 2]になります。