【LeetCode #75】Sort Colors
はじめに
今日の問題は medium レベルの問題です．
sort() を使えば簡単に解けますが，sort() 禁止のルールで解けという問題でした．
また，ワンパスで解ける？という煽りもありましたのでそちらのアプローチも掲載しています．（ワンパスについては下記参照）
具体的なアルゴリズムでは今回の問題は，カウントソート，オランダ国旗アルゴリズムというものを使用しています．
ワンパス
- 配列・文字列・入力を1回だけ前から後ろへ走査して完結する処理のこと
- 2 重ループやループ 2 つはダメ
問題概要
赤・白・青合計 n 個の要素で構成された配列に対して，in-place で sort をする問題です．
ここで，0 = 赤，1 = 白，2 = 青，を表します．
sort() を使えば一発で解けそうですが、この問題では built-in の sort 関数は禁止 されています．
そのため、自力でソートアルゴリズムを実装する必要があります．
例
入力: nums = [2, 0, 2, 1, 1, 0]
出力（in-place）: [0, 0, 1, 1, 2, 2]
カテゴリ
問題難易度：medium
タグ：Array，Two Pointers，Sorting
解法①：カウントソート
出現する数字が小さいときに有効な方法です．
具体的には，配列内に出現する値と出現回数を覚えておき，nums を小さい値から出現回数分ずつ埋めていきます．
今回は出現する値が 0，1，2 で固定されているので count が予め初期化できます．
具体例 1
nums = [2, 0, 2, 1, 1, 0]
- カウント
count = {0: 2, 1: 2, 2: 2} - 詰め直し
count = {0: 2, 1: 2, 2: 2} nums = [0, 0, 1, 1, 2, 2]
具体例 2
nums = [2, 0, 1]
- カウント
count = {0: 1, 1: 1, 2: 1} - 詰め直し
count = {0: 1, 1: 1, 2: 1} nums = [0, 1, 2]
実装はこちらです．
class Solution:
def sortColors(self, nums: List[int]) -> None:
"""
Do not return anything, modify nums in-place instead.
"""
count = {i: 0 for i in range(3)}
for n in nums:
count[n] += 1
k = 0
i = 0
while i < len(nums):
n = count[k]
for _ in range(n):
nums[i] = k
i += 1
k += 1
コード的には分かりやすいですし，こちらで問題なく AC しますが，ループが複数個ありワンパスではありません．
解法②：オランダ国旗アルゴリズム
こちらがワンパスで解く解法です．
left，mid，right の 3 ポインタを使います．それぞれの指す位置は下記です．
| ポインタ | 説明 |
|---|---|
| left | ソート済みの 0 の末尾の次の位置 |
| mid | 走査する値の位置 |
| right | ソート済みの 2 の先頭の前の位置 |
mid が指す値によって，下記の動作をします．
- mid == 0 のとき left の指す値と mid の指す値を交換し，left・mid どちらも +1 する
- mid == 1 のとき mid を +1 する
- mid == 2 のとき mid の指す値と right の指す値を交換し，right - 1 する（交換後の mid の値を再評価するため mid は進めない）
具体例 1
nums = [2, 0, 2, 1, 1, 0]
-
初期値
2, 0, 2, 1, 1, 0 ↑left ↑right mid -
1 回目
0, 0, 2, 1, 1, 2 ↑left ↑right mid -
2 回目
0, 0, 2, 1, 1, 2 ↑left ↑right mid -
3 回目
0, 0, 2, 1, 1, 2 ↑left ↑right mid -
3 回目
0, 0, 1, 1, 2, 2 ↑left ↑right mid -
4 回目
0, 0, 1, 1, 2, 2 ↑left ↑mid right -
5 回目（終了）
0, 0, 1, 1, 2, 2 ↑left ↑right ↑mid
具体例 2
nums = [2, 0, 1]
-
初期値
2, 0, 1 ↑left ↑right mid -
1 回目
1, 0, 2 ↑left ↑right mid -
2 回目
1, 0, 2 ↑left ↑mid right -
3 回目（終了）
0, 1, 2 ↑left ↑mid right
このような流れで 0，1，2 の並び替えが完了します．
ちなみに，オランダ国旗アルゴリズムはダイクストラ法で有名なエドガー・ダイクストラさんが考案したそうです．
実装はこちらです．
class Solution:
def sortColors(self, nums: List[int]) -> None:
"""
Do not return anything, modify nums in-place instead.
"""
left = 0
mid = 0
right = len(nums) - 1
while mid <= right:
if nums[mid] == 0:
nums[left], nums[mid] = nums[mid], nums[left]
left += 1
mid += 1
elif nums[mid] == 1:
mid += 1
elif nums[mid] == 2:
nums[mid], nums[right] = nums[right], nums[mid]
right -= 1
カウントソート vs オランダ国旗アルゴリズム
| カウントソート | オランダ国旗アルゴリズム | |
|---|---|---|
| パス数 | 2 パス | 1 パス |
| 空間計算量 | O(k)（今回の場合 O(1)） | O(1) |
| 時間計算量 | O(n+k)（今回の場合 O(n)） | O(n) |
| 対応する値の種類数 | 値が整数で数が少なければ任意に対応可 | 3種類のみ対応（赤・白・青）などに限定 |
TypeScript
解法①
/**
Do not return anything, modify nums in-place instead.
*/
function sortColors(nums: number[]): void {
let count = new Map<number, number>();
count.set(0, 0);
count.set(1, 0);
count.set(2, 0);
for (const n of nums){
const tmp: number = count.get(n) + 1;
count.set(n, tmp);
}
let k: number = 0;
for (let i: number = 0; i < 3; i++){
const c = count.get(i) || 0;
for (let j = 0; j < c; j++) {
nums[k++] = i;
}
}
};
解法②
/**
Do not return anything, modify nums in-place instead.
*/
const swap = (arr: number[], i: number, j: number): void => {
const tmp: number = arr[i];
arr[i] = arr[j];
arr[j] = tmp;
}
function sortColors(nums: number[]): void {
let left: number = 0;
let mid: number = 0;
let right: number = nums.length-1;
while (mid <= right){
switch (nums[mid]){
case 0:
swap(nums, left, mid);
left += 1;
mid += 1;
break;
case 1:
mid += 1;
break;
case 2:
swap(nums, mid, right);
right -= 1;
break;
}
}
};
まとめ
今回は built-in の sort を使用せず，0, 1, 2 の値を sort する方法をご紹介しました．
読みやすい・分かりやすいのはカウントソート，1 パスでとなるとオランダ国旗アルゴリズム，といった感じで使い分けると良いかなと思います！