Computer >> 컴퓨터 >  >> 프로그래밍 >> JavaScript

JavaScript로 정렬된 2차원 배열에서 N번째로 작은 요소 찾기

문제 상황

다음과 같이 각 행이 오름차순으로 정렬된 숫자 2차원 배열이 있다고 가정해 보겠습니다.

const arr = [
    [ 1, 5, 9],
    [10, 11, 13],
    [12, 13, 15]
];

우리가 작성해야 할 JavaScript 함수는 첫 번째 인수로 이러한 배열을, 두 번째 인수로 하나의 정수 num을 받습니다. 그리고 이 함수는 배열 arr에 존재하는 요소들 중 num번째로 작은 값을 반환해야 합니다.

예를 들어, 함수에 다음과 같이 입력을 전달한다면 −

const arr = [
    [ 1, 5, 9],
    [10, 11, 13],
    [12, 13, 15]
];
const num = 5;

출력

출력 결과는 다음과 같아야 합니다 −

const output = 11;

출력 설명

11은 이 행렬에서 다섯 번째로 작은 요소입니다. 실제로 배열의 모든 요소를 오름차순으로 나열하면 [1, 5, 9, 10, 11, 12, 13, 13, 15]가 되며, 다섯 번째 값이 바로 11입니다.

풀이 코드

이 문제는 이진 탐색(Binary Search)을 활용하여 효율적으로 해결할 수 있습니다. 코드는 다음과 같습니다 −

const arr = [
    [ 1, 5, 9],
    [10, 11, 13],
    [12, 13, 15]
];
const num = 5;
const kthSmallest = (arr = [], num = 1) => {
    let low = arr[0][0]
    let high = arr[arr.length-1][arr[0].length-1] + 1;
    while (low < high) {
        let mid = low + Math.floor((high-low)/2);
        let count = 0;
        for (let i = 0;i<arr.length;i++) {
            for (let j=0;j<arr.length;j++) {
                if (arr[i][j] <= mid) count++;
                else break;
            }
        }
    if (count < num) low = mid+1;
        else high = mid;
    }
    return low
};
console.log(kthSmallest(arr, num));

코드 설명

핵심 아이디어

일반적인 정렬된 1차원 배열에서 특정 값을 찾을 때는 인덱스 범위를 활용합니다. 예를 들면 다음과 같습니다.

low = 0, high = length-1, mid = (low+high)/2

찾으려는 값이 mid 인덱스의 값보다 크면 오른쪽 영역을 탐색하고, 작으면 왼쪽 영역을 탐색하는 방식입니다.

그러나 이 문제처럼 정렬된 2차원 배열에서는 그런 mid 인덱스를 직접 찾을 수 없습니다. 여기서의 핵심 발상은 인덱스 대신 값 자체의 범위를 기준으로 삼는 것입니다.

값의 범위를 활용한 이진 탐색

배열의 첫 번째 요소(arr[0][0])가 최솟값이고, 마지막 요소가 최댓값이라는 점을 이용하면, 우리가 찾는 답은 반드시 이 두 값 사이 어딘가에 존재합니다. 따라서 이 두 값을 low와 high로 설정하고, 그 사이의 중간값(mid)을 후보로 삼아 배열 내에서 mid보다 작거나 같은 숫자가 몇 개인지 세어 봅니다. 그 개수가 num보다 적으면 low를 올리고, 많거나 같으면 high를 내리는 방식으로 범위를 좁혀 나갑니다.

mid가 배열에 없는 값일 수 있는 이유와 해결 방법

여기서 매우 까다로운 부분이 등장합니다. mid를 계산하는 방식상, 우리가 검사하는 값이 실제 배열에는 존재하지 않는 임의의 숫자일 수 있다는 점입니다. 하지만 걱정할 필요가 없습니다.

low와 high가 서로 거의 만나게 되는 시점을 상상해 보면 이유를 알 수 있습니다. 만약 예상보다 적은 수의 숫자를 세었다면 low = mid + 1로 설정하게 되는데, 이는 사실상 mid를 1씩 증가시키는 것과 같습니다. 이렇게 mid를 한 칸씩 올려가다 보면 결국 반드시 배열에 실제로 존재하는 값에 도달하게 되고, 그 시점에 정확히 k개의 숫자를 세게 되므로 우리가 원하는 답을 얻을 수 있습니다.

실행 결과

콘솔에 출력되는 결과는 다음과 같습니다 −

11