[JS 알고리즘] Toy 18 - getItemFromTwoSortedArrays (kth element of two sorted arrays)

해당 문제는 코드스테이츠의 Toy problem을 출처로 합니다.

문제

길이가 m, n이고 오름차순으로 정렬되어 있는 자연수 배열들을 입력받아 전체 요소 중 k번째 요소를 리턴해야 합니다.

 

Advanced

• 단순히 처음부터 끝까지 찾아보는 방법(O(K)) 대신 다른 방법(O(logK))을 탐구해 보세요.

힌트
- 이진 탐색(binary search)을 응용하여 해결합니다.

입력

인자 1 : arr1

• 자연수를 요소로 갖는 배열
• arr1.length는 m

인자 2 : arr2

• 자연수를 요소로 갖는 배열
• arr2.length는 n

인자 3 : k

• number 타입의 0 이상의 정수

출력

• number 타입을 리턴해야 합니다.

주의사항

• 두 배열의 길이의 합은 1,000,000 이하입니다.
• 어떤 배열 arr의 k번째 요소는 arr[k-1]을 의미합니다.

입출력 예시

let arr1 = [1, 4, 8, 10];
let arr2 = [2, 3, 5, 9];
let result = getItemFromTwoSortedArrays(arr1, arr2, 6);
console.log(result); // --> 8

arr1 = [1, 1, 2, 10];
arr2 = [3, 3];
result = getItemFromTwoSortedArrays(arr1, arr2, 4);
console.log(result); // --> 3

 

---

문제풀이

 

원리

두 배열 모두 오름차순으로 정렬되어있다는 사실에 초점을 맞춰보자.

 

오름차순으로 정렬된 배열을 반으로 나눌 경우, 오른쪽에 위치한 요소는 왼쪽에 위치한 요소보다 언제나 더 큰 값을 가진다.

 오름차순으로 정렬된 배열을 반으로 쪼개 left array와 right array 둘로 나눈다고 생각해보자. 그리고 k번째 요소가 left array의 마지막에 위치할 것이라고 가정해보자. 그렇다면 k번째 요소가 left array에 위치한 이상, right array는 반드시 k번째 요소보다 큰 값들만 존재한다. 배열은 오름차순으로 정렬되어있기 때문에, right array의 모든 요소들은 left array의 어떤 요소들이든 보다 항상 크다. 

 위 이미지의 예시로 원칙이 맞는지 확인해보자. 우리가 찾아야하는 k번째 요소는 6이라고 하자.

그리고 입력받은 array1은 [2,3,6,7,9]이고 array2는 [1,4,8,10] 이라고 하자.

 

 인덱스 위치를 찾아야하는 값은 6이며, 우리는 배열을 둘로 쪼개 왼쪽 배열(left array)의 끝부분에 찾아야하는 값이 있어야한다고 조건을 걸어놨다.

 따라서 6을 기준으로 반을 쪼개면, array1의 left array는 [2,3,6]이고 right array는 [7,9]가 될 것이다. 6은 right array의 7보다 작고, 당연히 9보다도 작다. 만약 array1이 left array는 [2,3]이고 right array는 [6,7,9]로 쪼개져있어도 마찬가지이다. left array의 끝값인 3은 right array의 요소 6,7,9보다 작다. 우리는 left array의 어떤 요소이든지 간에 right array의 모든 요소보다 작은 값만 모여있다는 것을 확인할 수 있다. 이는 당연히 array2도 마찬가지로 적용된다.

 

그러므로, 지금 찾은 원리와 함께 이진 탐색(binary search)를 활용하여 계속 반으로 쪼개가며 left half array에 k번째 요소가 있는지를 찾을 것이다.

 

 

접근법

두 배열을 각각 array1, array2라고 한다. 그리고 반으로 쪼개가며 탐색하기 위한 left 포인터와 right 포인터를 만들고, 이를 가지고 중앙값(middle)을 찾는다. 그리고 k번째 값이 left array 마지막에 위치하는 인덱스를 찾기 위해 아래와 같은 변수를 만들어준다.

변수	설명
l1	array1을 반으로 나눠 얻은 left array의 마지막 인덱스 값
r1	array1을 반으로 나눠 얻은 right array의 첫번째 인덱스 값
l2	array2을 반으로 나눠 얻은 left array의 마지막 인덱스 값
r2	array2을 반으로 나눠 얻은 right array의 첫번째 인덱스 값

l1 <= r2  AND  l2 <= r1

찾고자 하는 k번째 요소가 왼쪽 배열에 있을 경우, 위 공식이 성립함을 계속 기억하고 활용해야한다. left half array에 위치한 요소들은 언제나 right half array의 요소들보다 작기때문이다.

 

 

이제 원리를 알았으니 l1, l2를 어떤 기준으로 구할 것인지가 관건이다. 이를 위한 변수를 추가로 선언해준다.

변수	설명
cut1	array1에서 가져올 요소의 수: (low+high)/2
cut2	array2에서 가져올 요소의 수: k-cut1
low	array1의 최소값 (array1에서 고를 수 있는 요소의 수 중에서 최소값)
high	array1의 최댓값 (array1에서 고를 수 있는 요소의 수 중에서 최댓값)

입력받은 두 배열을 합쳐 오름차순으로 정렬한 배열을 k번째 요소를 기준으로 나눴을 때의 왼쪽 배열의 개수는 k개이다.

 array1과 array2를 합쳐 오름차순으로 정렬한 배열을 Arr라고 해보자.

우리는 앞서 찾고자하는 k번째 요소가 Arr의 왼쪽 배열 마지막 인덱스에 위치하게끔 반으로 나눌 것이라고 했다.

이는 Arr의 왼쪽 배열의 개수가 k개라는 뜻이기도 하다. 이를 이용해 두 배열에서 가져올 요소의 수를 정할 수 있다.

cut1이 중앙값을 기준으로 array1에서 가져올 요소의 개수를 정하면, cut2에서 가져올 수 있는 요소의 수는 k를 넘을 수 없기에 k-cut1이 된다. array1의 왼쪽배열과 array2의 왼쪽배열의 요소 개수의 합은 Arr의 왼쪽배열 요소의 개수와 같기 때문이다.

 

 

보다 자세한 접근법은 정리해둔 이미지로 대체하겠다.

 

작성 코드 (Javascript)

const getItemFromTwoSortedArrays = function (arr1, arr2, k) {
  let m = arr1.length;
  let n = arr2.length;
  // m < n 이 되도록 reverse 시켜준다 (가장 작은 것에 이진탐색하기 위함)
  // 이 가정을 성립시켜놔야 밑에 low=min(0, k-n) 로 초기화할 수 있음
  if (m > n) {
    return getItemFromTwoSortedArrays(arr2, arr1, k);
  }
  // low, high 초기화
  let low = Math.max(0, k-n);
  let high = Math.min(k, m);

  // low <= high인 동안 값을 찾는다
  while (low <= high) {
    let cut1 = (low + high) >> 1;
    let cut2 = k - cut1;
    let l1 = cut1 === 0 ? Number.MIN_SAFE_INTEGER : arr1[cut1 - 1];
    let l2 = cut2 === 0 ? Number.MIN_SAFE_INTEGER : arr2[cut2 - 1];
    let r1 = cut1 === m ? Number.MAX_SAFE_INTEGER : arr1[cut1];
    let r2 = cut2 === n ? Number.MAX_SAFE_INTEGER : arr2[cut2];
    // l1 <= r2 && l2 <= r1
    if (l1 <= r2 && l2 <= r1) {
      return Math.max(l1, l2);
    }
    // l1 > r2
    if (l1 > r2) {
      high = cut1 - 1;
    }
    // l2 > r1
    else {
      low = cut1 + 1;
    }
  }
  return -1;
};

 

 

코드스테이츠 레퍼런스 코드 (Javascript)

// naive solution --- O(n) solution
// const getItemFromTwoSortedArrays = function (arr1, arr2, k) {
//   let cnt = 0,
//     left = 0,
//     right = 0;
//   let target;
//   while (cnt < k) {
//     if (arr1[left] < arr2[right]) {
//       target = arr1[left];
//       left++;
//     } else {
//       target = arr2[right];
//       right++;
//     }
//     cnt++;
//   }
//   return target;
// };

// O(logK) solution
const getItemFromTwoSortedArrays = function (arr1, arr2, k) {
  let leftIdx = 0,
    rightIdx = 0;

  // k번째 요소를 찾는 것 === 앞에서부터 k번 count했을 때의 요소
  // 이진 탐색법을 이용하기 위해 주어진 두 개의 배열(arr1, arr2)를 합치지 않고, 
  // k를 절반으로 나누어 각각의 배열에서 따로 count를 셉니다.
  while (k > 0) {
    // 이진 탐색을 위해 각 배열에서 k를 절반으로 쪼개서 카운트 한다.
    let cnt = Math.ceil(k / 2);
		// leftStep은 leftArr(arr1)의 count 할당량(k/2)
		// rightStep은 rightArr(arr2)의 count 할당량(k/2)이라고 생각합니다.
    let leftStep = cnt,
      rightStep = cnt;

    // 엣지 케이스
    // 할당된 카운트가 남았음에도 배열의 끝에 도달하면 k를 나머지 배열 쪽으로 넘긴다.
    if (leftIdx === arr1.length) {
      rightIdx = rightIdx + k;
      break;
    }

    if (rightIdx === arr2.length) {
      leftIdx = leftIdx + k;
      break;
    }

    // 엣지 케이스
    // 현재 카운트가 남아있는 후보 요소들보다 많을 경우, 
		// leftStep(현재 할당량)을 남아있는 요소들의 개수로 바꾼다.
    if (cnt > arr1.length - leftIdx) leftStep = arr1.length - leftIdx;
    if (cnt > arr2.length - rightIdx) rightStep = arr2.length - rightIdx;

		// 본격적으로 카운트하는 로직
    // 두 배열의 현재 검사 요소 위치를 비교해서, 그 값이 작은 배열은 비교한 위치 앞에 있는 요소들을 모두 후보군에서 제외시킨다.
// TODO 이진 탐색으로 가기 때문에, 두 배열의 현재 cursor(Idx)에서 현재 진행해야하는 count수(leftStep/rightStep)를 합친 index의 값을 비교합니다.
// * (가장 첫 시행 때는, k가 100인 경우 두 배열의 50번째 Idx 값을 비교하게 되겠습니다. Idx는 0, Step은 50입니다.)
// * 비교 후 그 값이 작은 배열은, 해당 값 앞의 요소에 대해서 다시 검사할 필요가 없습니다. count에 포함해도 되는, 비교된 큰 값보다 무조건 작은 값이기 때문입니다.
// * 그러므로, 값이 작은 쪽의 배열의 cursor(Idx)를 옮겨줍니다. 다음 turn 부터는 cursor 이후로만 탐색합니다. 이러한 동작을 이진 탐색이라고 할 수 있겠습니다.
// ! Right가 더 큰 경우, 작은 값은 무시된다.
    if (arr1[leftIdx + leftStep - 1] < arr2[rightIdx + rightStep - 1]) {
      leftIdx = leftIdx + leftStep;
      // 처리가 끝나면 k값을 절반으로 떨어뜨린다.
// * (위의 "count에 포함해도 되는 무조건 작은 값이기 때문입니다."에 의해서, 후보군에서 떨어진 요소들은 무시하게되는 처리입니다.)
      k = k - leftStep;
    } else { // ! Left가 더 큰 경우, 작은 값은 무시된다.
      rightIdx = rightIdx + rightStep;
      k = k - rightStep;
    }
  }
// while문을 끝내고 나왔을 때, 두 개의 배열은 조정된 cursor를 가질 것이고, 해당 값 중 더 큰 값이 k번째 요소일 것입니다. (k번째를 벗어나지 않습니다.)
  leftMax = arr1[leftIdx - 1] || -1;
  rightMax = arr2[rightIdx - 1] || -1;

  return Math.max(leftMax, rightMax);
};

 

 

---

 

참고 문헌 : https://takeuforward.org/data-structure/k-th-element-of-two-sorted-arrays/

참고 강의 : https://youtu.be/nv7F4PiLUzo