[LeetCode] 378. Kth Smallest Element in a Sorted Matrix

문제 출처: https://leetcode.com/problems/kth-smallest-element-in-a-sorted-matrix/

Kth Smallest Element in a Sorted Matrix - LeetCode

 

문제

각 행과 열이 오름차순으로 정렬되어 있는 n x n 크기의 matrix가 주어지면, 행렬 내에서 k번째로 작은 원소를 반환하라.

 

정렬된 순서에서 k번째로 작은 원소를 찾는 것이지, k번째로 구분되는 원소를 찾는 것이 아니다.

 

예제

Input: matrix = [[1,5,9],[10,11,13],[12,13,15]], k = 8
Output: 13

Explanation: The elements in the matrix are [1,5,9,10,11,12,13,13,15], and the 8th smallest number is 13

Input: matrix = [[-5]], k = 1
Output: -5

 

풀이

가장 먼저 떠올린 방법은 "2차원을 1차원으로 변형시킨 후, 정렬해서 인덱스로 접근하자."입니다. 단순하지만 확실하게 k번째로 작은 값에 도달할 수 있는 방법입니다. 코드도 단순하게 작성되어집니다.

var kthSmallest = function(matrix, k) {
    const elements = matrix.flat();
    elements.sort((a, b) => a - b);

    return elements[k-1];
}

이러한 방법에서 시간 복잡도는 정렬에 걸리는 시간에 의존하게 됩니다. 보통 정렬 알고리즘의 경우 O(nlogn)의 시간 복잡도를 갖기 때문에 이 문제의 경우 matrix의 크기가 n x n이므로, 이 행렬의 원소는 모두 n^2개가 되며 O(n^2 x logn^2)의 시간이 걸릴 것이라 예상할 수 있습니다.

 

이를 개선하기 위해서 문제에서 언급은 되었던 행과 열이 정렬되어 있다는 조건을 활용할 수 있다.

 

정렬되어있는 배열에서 특정한 값을 찾는 알고리즘.... 이진 탐색 알고리즘을 적용할 수 있을... 까?

 

이진 탐색 알고리즘을 적용하려고 보니 문제가 주어지는 배열이 2차원 배열이라는 점이다. 그래서 이 부분에 대해서 고려하지 않고 1차원 정렬된 배열이 있다고 가정하고 배열에서 k번째로 작은 값을 찾는 방법을 이진탐색 알고리즘을 적용해 찾는 방법에 대해 생각해보자.

// 만약 주어지는 값이 1차원 정렬된 리스트라면?
var kthSmallest = function(arr, k) {
    let l = 0, r = arr.length-1;
    while (l <= r) {
        const m = l + Math.floor((r-l)/2);

        if ((m+1) < k) {
            l = m + 1;
        } else {
            r = m - 1;
        }
    }
    return l;
}

이진 탐색을 위하여 배열의 시작 인덱스와 마지막 인덱스의 위치를 각각 l과 r로 설정한다. 이 두 변수로 구해지는 중간값인 m은 말 그대로 중간 위치를 가리키는 값이 된다. 이 m값은 인덱스이기도 하지만, 몇 번째로 작은 값인지에 대한 정보를 담고 있는 값이기도 하다. 왜냐하면 배열이 정렬된 상태이기 때문이다.

 

이 m 값을 k와 비교하여, 만약 m < k라면, l를 이동시켜 이후 계산될 m이 더욱 커지게 하고, m >= k면, r을 이동시켜 이후 m값을 더욱 작게 하는 방향으로 찾고자 하는 값을 찾아나갈 수 있다.

 

실제 문제 상황인 2차원 배열 상태일 때를 생각해보자. 2차원 배열은 1차원 배열과 다르게 중간 인덱스를 특정하기 까다롭다. 고려해야 되는 변수가 두 개(row, column)이기 때문이다. 만약 위치를 특정한다고 하더라도, 그 위치를 토대로 해당 위치의 값이 몇 번째로 작은 수인지 알아내기는 어렵다.

 

그래서, 약간의 트릭(?)을 사용해서 문제를 접근해볼 수 있다.

 

기존에는 이진 탐색을 위한 변수 l과 r 그리고 m이 배열의 인덱스를 가리켰다면, l과 r이 값을 가리키도록 하는 것이다. l은 matrix내 가장 작은 값(matrix[0][0]), r은 matrix내 가장 큰 값(matrix[n-1][n-1])을 갖게 하는 것이다. 이와 같은 방식으로 접근하면 m 의 경우 l과 r의 중간값일 것이고, matrix에 존재하는 값일 수도 존재하지 않는 값일 수도 있다.

 

1차원 배열이라 가정했을 때는 m과 k값을 직접 비교하며 l과 r을 설정해줬는데, 그 이유는 m이 배열 내에서 몇 번째로 작은 값인지에 대한 정보를 갖고 있었기 때문이었다. 마찬가지로 2차원 배열에서도 현재 m보다 작은 값이 몇 개 있는지 알려주는 lessOrEqualCount() 함수를 만들어 볼 수 있고 이 함수가 있다고 한다면 아래와 같은 코드로 문제를 해결할 수 있다.

var kthSmallest = function(matrix, k) {
    const n = matrix.length;
    let l = matrix[0][0], r = matrix[n-1][n-1];
    while (l <= r) {
        const m = l + Math.floor((r-l)/2);

        const count = lessOrEqualCount(matrix, m)    // TODO

        if (count < k) {
            l = m + 1;
        } else {
            r = m - 1;
        }
    }
    return l;
}

이제 lessOrEqualCount(matrix, m)라는 함수를 구현해보자. 주어지는 행렬이 행과 열에 대해 오름차순 정렬되어있다는 사실을 이용할 수 있다.

 

만약 m = 10이고, matrix[i][j] = 15라고 하면, matrix[i+1 ~ n-1][j+1 ~ n-1]에 속하는 값들은 확인하지 않아도 된다. 오름차순 정렬되어 있기에 무조건 15보다 크거나 같은 값일 것이기 때문이다. 이 점을 이용하여 다음과 같은 방식으로 m보다 작은 원소의 수를 찾아낼 수 있다.

var lessOrEqualCount(matrix, target) {
    const n = matrix.length;

    let count = 0;
    let j = n-1;

    for (let i = 0; i < n; i++) {
        while (j >= 0 && matrix[i][j] > target) { j -= 1; }
        count += j+1;
    }

    return count;
}

 

Github: 378.js

opwe37/Algorithm-Study