세 구역 분할

언제 쓰나

배열의 원소를 작은 무리, 가운데 무리, 큰 무리 세 무리로 나눠, 앞에서부터 그 순서로 모아야 할 때 씁니다.

• 값이 세 종류뿐인 배열을 정렬합니다. 이 글의 대표 문제처럼 0, 1, 2가 섞인 배열을 0, 1, 2 순서로 모으는 경우입니다.
• 기준값 하나를 두고 “기준보다 작은 수 / 같은 수 / 큰 수”로 나눕니다. 퀵 정렬·퀵셀렉트에서 같은 수가 많은 입력이 느려질 때, 분할을 이 방법으로 바꿉니다.
• 새 배열을 만들지 않고 입력 배열 안에서, 배열을 한 번만 훑어서 끝내야 합니다.

어떻게 동작하나

개수를 세면 안 되나

LeetCode 75는 0, 1, 2가 섞인 배열을 라이브러리 정렬 함수 없이 0, 1, 2 순서로 정렬하는 문제입니다. nums = [2, 0, 2, 1, 1, 0]이면 [0, 0, 1, 1, 2, 2]가 되어야 합니다.

이 문제만 보면 개수를 세는 방법으로도 충분합니다. 한 번 훑어 0이 2개, 1이 2개, 2가 2개라는 것을 센 뒤, 앞에서부터 0을 두 칸, 1을 두 칸, 2를 두 칸 다시 적으면 됩니다. 시간도 O(n)입니다.

개수 세기가 통하는 것은 같은 무리의 값이 모두 같기 때문입니다. 0 무리는 전부 0이라 개수만 알면 다시 적을 수 있습니다. 그런데 [3, 9, 1, 5, 5, 7]을 기준값 5로 나눈다면, 5보다 작은 무리는 3과 1, 큰 무리는 9와 7입니다. “작은 수 2개”라는 개수만으로는 3과 1을 되살릴 수 없으니, 원소를 실제로 자리에서 옮겨야 합니다.

세 구역 분할은 원소를 바꿔 넣으며 제자리에서 세 무리로 나누는 방법입니다. 그래서 0·1·2 정렬뿐 아니라 기준값으로 나누는 경우에도 그대로 쓸 수 있고, 배열을 한 번만 훑습니다. LeetCode 75도 후속 질문으로 바로 이 “한 번만 훑고 추가 공간은 변수 몇 개만 쓰는 방법”을 묻습니다.

0은 앞으로, 2는 뒤로 보낸다

기본 발상은 이렇습니다. 배열을 왼쪽부터 한 칸씩 읽으면서 0은 배열 앞쪽으로, 2는 배열 뒤쪽으로 보내고, 1은 그 자리에 둡니다. 앞쪽에는 0이 왼쪽부터 차곡차곡 쌓이고, 뒤쪽에는 2가 오른쪽부터 쌓이니, 다 읽고 나면 그 사이에 1만 남습니다.

이렇게 하려면 위치 세 개를 기억해야 하고, 코드에서는 이 셋을 변수 lo, mid, hi에 담습니다. 이렇게 배열의 칸 번호를 담아 가리키는 변수를 포인터(pointer)라고 합니다.

• lo: 다음 0을 놓을 자리입니다. 앞쪽에 쌓인 0들의 바로 뒤 칸이고, 처음에는 맨 앞 칸입니다.
• hi: 다음 2를 놓을 자리입니다. 뒤쪽에 쌓인 2들의 바로 앞 칸이고, 처음에는 맨 뒤 칸입니다.
• mid: 지금 읽는 칸입니다. 이름 때문에 배열의 한가운데를 가리킬 것 같지만, 실제로는 1이 모인 가운데 무리의 바로 뒤 칸을 가리키고, 그 칸이 곧 다음에 읽을 칸입니다.

mid 칸의 값을 읽고 할 일은 셋 중 하나입니다.

1. 0이면 lo 칸과 값을 바꿔 0을 앞으로 보냅니다. 0이 한 칸 쌓였으니 lo를 한 칸 오른쪽으로 옮기고, mid도 다음 칸으로 갑니다.
2. 1이면 그대로 두고 mid만 다음 칸으로 갑니다.
3. 2이면 hi 칸과 값을 바꿔 2를 뒤로 보냅니다. 2가 한 칸 쌓였으니 hi를 한 칸 왼쪽으로 옮깁니다. mid는 그대로 둡니다.

[2, 0, 1]로 직접 해 보겠습니다. 처음에 lo와 mid는 0번 칸, hi는 2번 칸입니다.

1. mid가 0번 칸의 2를 읽습니다. hi 칸(2번)의 1과 바꾸면 [1, 0, 2]가 되고, 맨 뒤의 2는 자리를 찾았으니 hi가 1번 칸으로 옵니다. mid는 0번 칸에 그대로 있습니다.
2. mid가 0번 칸을 다시 읽습니다. 방금 뒤에서 넘어온 1이라 그대로 두고, mid가 1번 칸으로 갑니다.
3. mid가 1번 칸의 0을 읽습니다. lo 칸(0번)의 1과 바꾸면 [0, 1, 2]가 되고, lo와 mid가 한 칸씩 오른쪽으로 갑니다.
4. 이제 mid(2번 칸)가 hi(1번 칸)를 지나쳤습니다. 읽지 않은 칸이 없으니 끝이고, 배열은 [0, 1, 2]로 정렬됐습니다.

세 칸을 정렬하는 동안 값을 읽은 횟수는 세 번이고, 배열을 앞에서 뒤로 한 번 지나간 것이 전부입니다.

왜 2와 바꿀 때만 mid가 멈추나

위 1번 과정에서 mid가 멈춘 이유는, 뒤에서 넘어온 1이 아직 아무도 읽지 않은 값이기 때문입니다. hi 칸은 읽기가 아직 닿지 않은 뒤쪽 칸이라, 그 자리에서 무엇이 넘어올지 모릅니다. 0이 넘어왔을 수도 있으니, mid는 그 자리에 남아 넘어온 값을 한 번 더 읽어야 합니다.

0을 lo 칸과 바꿀 때는 사정이 다릅니다. lo는 늘 mid와 같은 칸이거나 그보다 앞에 있어서, lo 칸의 값은 mid가 이미 지나오며 읽은 값입니다. 0은 이미 앞쪽에 쌓아 두었으니 그 자리에 남아 있는 값은 1뿐이고, 1은 가운데에 두면 되므로 mid가 바로 다음 칸으로 가도 됩니다. lo와 mid가 같은 칸이면 자기 자신과 바꾸는 것이라 값이 달라지지 않습니다.

읽는 도중의 배열은 네 구역으로 나뉜다

규칙대로 읽다 보면 배열은 언제나 네 구역으로 나뉩니다. 아래는 더 긴 배열을 정렬하는 도중의 한 장면입니다.

네 구역 · nums = [0, 1, 1, 2, 0, 1, 2, 2]

• lo 앞은 0 구역입니다. 앞으로 보낸 0이 쌓인 곳입니다.
• lo부터 mid 바로 앞까지는 1 구역입니다. 읽고 그대로 둔 1이 모인 곳입니다.
• mid부터 hi까지는 안 본 칸입니다. 노란 칸이 지금 읽는 mid 칸입니다.
• hi 뒤는 2 구역입니다. 뒤로 보낸 2가 쌓인 곳입니다.

한 번 읽을 때마다 mid가 오른쪽으로 가거나 hi가 왼쪽으로 오니, 안 본 칸은 한 칸씩 줄어듭니다. 안 본 칸이 없어지면 0 구역, 1 구역, 2 구역만 남아 정렬이 끝납니다. 이 문제는 에츠허르 데이크스트라(Edsger Dijkstra)가 네덜란드 국기의 세 색에 빗대어 낸 문제라 네덜란드 국기 문제(Dutch national flag problem)라고도 부릅니다.

직접 따라가 보기

아래에서 “다음”을 누르면 mid가 한 칸을 읽을 때마다 포인터와 구역 괄호가 함께 움직입니다. 2번째 단계에서 맨 앞의 2가 맨 뒤의 0과 교환 후보로 이어지고, 3번째 단계에서 두 값이 바뀝니다. 이때 hi만 왼쪽으로 한 칸 오고 mid는 맨 앞에 그대로 남는 것을 보세요. 4번째 단계에서 mid가 그 자리로 온 0을 읽고 나서야 0 구역이 생깁니다. 6~7번째 단계에서도 같은 장면이 한 번 더 나옵니다.

처음부터 이전 자동 재생 다음

단계 1 / 10

세 구역 분할 · nums = [2, 0, 2, 1, 1, 0]

변수

lo

0다음 0이 들어갈 칸

mid

0지금 읽는 칸

hi

5다음 2가 들어갈 칸

nums

–

lo와 mid는 맨 앞, hi는 맨 뒤에서 시작합니다. 여섯 칸 모두 아직 안 본 칸입니다. "다음"을 눌러 보세요

• 지금 읽는 칸 (mid)
• 자리가 정해진 칸 (0 구역·2 구역)
• 가운데에 모인 1
• 아직 안 본 칸

대표 문제 풀이

아래는 참고 풀이입니다. 반복문 안의 ①②③이 위의 세 규칙과 같은 순서이고, 변수 이름은 시각화의 변수 칸과 같습니다.

function sortColors(nums) {
  let lo = 0;
  let mid = 0;
  let hi = nums.length - 1;

  while (mid <= hi) {
    if (nums[mid] === 0) {
      // ① 0은 앞으로: lo 칸과 바꾸고 둘 다 오른쪽으로
      [nums[lo], nums[mid]] = [nums[mid], nums[lo]];
      lo++;
      mid++;
    } else if (nums[mid] === 1) {
      // ② 1은 그대로: mid만 오른쪽으로
      mid++;
    } else {
      // ③ 2는 뒤로: hi 칸과 바꾸고 hi만 왼쪽으로. mid는 그대로
      [nums[mid], nums[hi]] = [nums[hi], nums[mid]];
      hi--;
    }
  }
}

반복 조건이 mid <= hi인 것은 mid와 hi가 같은 칸일 때도 안 본 칸이 하나 남아 있기 때문입니다. 반복 한 번마다 mid가 오른쪽으로 가거나 hi가 왼쪽으로 와서 안 본 칸이 한 칸씩 줄어드니, 반복은 최대 n번이고 시간은 O(n)입니다. 입력 배열 안에서 값만 바꾸고 변수 세 개만 더 쓰므로 공간은 O(1)입니다.

자주 하는 실수

2를 hi 칸과 바꾼 뒤에도 mid를 옮깁니다. 세 경우 모두 mid++를 넣으면 코드 모양은 더 고르지만, hi 칸에서 온 값을 읽지 않고 지나치게 됩니다. [1, 2, 0]에서는 2와 맨 뒤의 0을 바꾼 뒤 mid가 그 0을 건너뛰어, 0이 1 구역에 남은 채 반복이 끝납니다.

hi와 바꾼 뒤 mid도 옮기면 · nums = [1, 2, 0]

이 실수는 LeetCode 75의 예제 1·2에서는 우연히 맞는 답을 내므로, 예제만 돌려 보고는 알아차리기 어렵습니다. [1, 2, 0]처럼 2 뒤에서 0이 넘어오는 입력으로 한 번 더 확인하면 드러납니다.

반복 조건을 mid < hi로 씁니다. mid와 hi가 같은 칸을 가리킬 때 그 칸은 아직 읽지 않은 칸인데, 조건이 <이면 이 칸을 읽지 않고 끝납니다. 예제 2의 [2, 0, 1]을 넣으면 [1, 0, 2]가 나옵니다.