1439. Find the Kth Smallest Sum of a Matrix With Sorted Rows

문제

m x n 2차원 배열이 있을때 각 행에서 1개씩 뽑아서 더한 값의 경우의 수 중 K번째로 작은 값을 구하기. 참고로 각 row는 오름차순 정렬되어있다.

Input: mat = [[1,3,11],[2,4,6]], k = 5
Output: 7
Explanation: Choosing one element from each row, the first k smallest sum are:
[1,2], [1,4], [3,2], [3,4], [1,6]. Where the 5th sum is 7.

https://leetcode.com/problems/find-the-kth-smallest-sum-of-a-matrix-with-sorted-rows/

해결 #1은 답은 맞지만 성능이 매우 떨어지는 풀이 방법. #2는 더 성능이 뛰어난 풀이방법. 따라서 정답은 해결 #2를 보면 됨.

해결 #1

// (8) 7 6 5 5 3 1   (8) is k'st smallest value in k-size max heap
void kth_smallest_add(int arr[], int val)
{
    if (q->heap_size < kth) {
        heap_push(arr, val);
    } else if (val < arr[0]) {
        arr[0] = val;
        heapify_down(arr, 0, q->heap_size);
    }
}

K 사이즈의 Max heap을 만들고 K개까지만 남도록 push하면 heap[0] 첫번째 인덱스가 K번째로 작은 값이 된다. 비슷한 해결방식은 (https://velog.io/@soopsaram/Leetcode-703.-Kth-Largest-Element-in-a-Stream) 여기서 풀었었다.

void back_tracking(int **mat, int row, int sum)
{
    /* base case */
    if (row >= row_size) { // 
        kth_smallest_add(q->heap, sum);
        return;   
    } 
    /* recursion */
    for (int col = 0; col < col_size; col++)
        back_tracking(mat, row + 1, sum + mat[row][col]);
}

2차원 배열에서 각 행(row)를 한개씩 선택하는 모든 경우의수 탐색하는 back tracking 함수 구조. 다만 모든 경우의 수를 back tracking 방식으로 (+가지치기) 탐색했는데, 29/72 에서 time out이 발생했다.

모든 TC패스를 못해서 완전한 풀이는 아니다. 백트레킹 자체가 O(m^n) 이기 때문에 엄청느린건 당연한듯 싶고 다른 방식으로 풀어야할듯! 가지치기를 잘해서 시간복잡도를 드라마틱하게 줄일수있을지는 잘 모르겠다

struct pq {
    int heap_size;
    int *heap;
};
struct pq *q;
int row_size;
int col_size;
int kth;

void swap(int arr[], int a, int b)
{
    int tmp = arr[a];
    arr[a] = arr[b];
    arr[b] = tmp;
}

void heapify_down(int arr[], int p_idx, int size)
{
    int largest = p_idx;
    int l_idx = p_idx * 2 + 1;
    int r_idx = p_idx * 2 + 2;
    
    if (l_idx < size && arr[l_idx] > arr[largest])
        largest = l_idx;
    if (r_idx < size && arr[r_idx] > arr[largest])
        largest = r_idx;
    if (largest != p_idx) {
        //swap(arr, largest, p_idx);
        int tmp = arr[largest];
        arr[largest] = arr[p_idx];
        arr[p_idx] = tmp;
        heapify_down(arr, largest, size);
    }
}

void build_heap(int arr[], int size)
{
    for (int i = (size  / 2) - 1; i >= 0; i--)
        heapify_down(arr, i, size);
}

void heapify_up(int arr[], int cur_idx)
{
    int p_idx = (cur_idx - 1) >> 1;
    
    if (cur_idx < 1)
        return;
    if (arr[cur_idx] > arr[p_idx]) {
        swap(arr, cur_idx, p_idx);
        heapify_up(arr, p_idx);
    }
}
void heap_push(int arr[], int val)
{
    arr[q->heap_size] = val;
    q->heap_size++;
    heapify_up(arr, q->heap_size - 1);
}

int heap_pop(int arr[])
{
    int ret = arr[0];
    q->heap_size--;
    arr[0] = arr[q->heap_size];
    heapify_down(arr, 0, q->heap_size);
    return ret;
}

// (8) 7 6 5 5 3 1   (8) is k'st smallest value in k-size max heap
void kth_smallest_add(int arr[], int val)
{
    if (q->heap_size < kth) {
        heap_push(arr, val);
    } else if (val < arr[0]) {
        arr[0] = val;
        heapify_down(arr, 0, q->heap_size);
    }
}

void back_tracking(int **mat, int row, int sum)
{
    /* TODO: need to pruning */
    if (q->heap_size > kth && sum > q->heap[0])
        return; 
    /* base case */
    if (row >= row_size) {
        kth_smallest_add(q->heap, sum);
        return;   
    } 
    /* recursion */
    for (int col = 0; col < col_size; col++) {
        back_tracking(mat, row + 1, sum + mat[row][col]);
    }
    return 0;
}

void print_arr(int arr[], int size)
{
    for (int i = 0; i < size; i++)
        printf("(%d)", arr[i]);
    printf("\n");
}

int kthSmallest(int** mat, int matSize, int* matColSize, int k){
    row_size = matSize;
    col_size = *matColSize;
    kth = k;
    int ret = 0;
    
    q = (struct pq *)malloc(sizeof(struct pq));
    memset(q, 0, sizeof(struct pq));
    q->heap = (int *)malloc(sizeof(int) * kth);
    memset(q->heap, 0, sizeof(int) * kth);
    
    back_tracking(mat, 0, 0);
    return q->heap[0];
}

해결 #2

제대로된 해결방법. 위의 방법은 잘못된 접근방법. 최종 결과는

Runtime: 38 ms, faster than 100.00% of C 
Memory Usage: 15.7 MB, less than 100.00% of C

m x n크기의 mat[][] 에서 딱 두개의 row씩만 모든 합의 경우의수를 구하고. 그것을 row갯수만큼 반복하면 모든 경우의 수의 sum값 배열을 얻을수 있을것이다(그런데 배열의 size는 n^m 개가 될것)

가령 mat = [[1,10,10],[1,4,5],[2,3,6]], k = 7 일때, 상위 두개의 row를 계산하면 [1+1, 1+4, 1+5, 10+1, 10+4, 10+5, 10+1, 10+4, 10+5] 가 될것이다. 그다음 루프에서 이 배열과 [2,3,6]을 같은 과정을 구하고 sort하면 최종 k번째 작은 sum값을 얻을 수 있다.

여기서 문제는 배열을 계속 계산하다보면 사이즈가 n^m 개가 되어버리는것. 그런데 다시 생각해보면 sum이 계산되는 배열의 최대 크기는 정렬된상태에서 작은순서대로 k개만 되면 된다. 따라서 배열을 heap구조로 만들고 k개만큼만 pop하면 해결할 수 있다.

시간 복잡도는 배열 계산 O(m * n * k)
heap 생성 O(m * k)
heap pop O(m * k * log k)

struct pq {
    int heap_size;
    int *heap;
};
struct pq *q;

void heapify_down(int arr[], int p_idx, int size)
{
    int min_idx = p_idx;
    int l_idx = p_idx * 2 + 1;
    int r_idx = p_idx * 2 + 2;
    
    if (l_idx < size && arr[l_idx] < arr[min_idx])
        min_idx = l_idx;
    if (r_idx < size && arr[r_idx] < arr[min_idx])
        min_idx = r_idx;
    if (min_idx != p_idx) {
        /* swap(arr, min_idx, p_idx); */
        int tmp = arr[min_idx];
        arr[min_idx] = arr[p_idx];
        arr[p_idx] = tmp;
        heapify_down(arr, min_idx, size);
    }
}

void build_heap(int arr[], int size)
{
    for (int i = (size  / 2) - 1; i >= 0; i--)
        heapify_down(arr, i, size);
}

int heap_pop(int arr[])
{
    int ret = arr[0];
    q->heap_size--;
    arr[0] = arr[q->heap_size];
    heapify_down(arr, 0, q->heap_size);
    return ret;
}
// mat[m][n]
int kthSmallest(int** mat, int matSize, int* matColSize, int k) {
    int res_size = *matColSize;
    int res_array_size = res_size > k ? res_size : k;
    int *res = (int *)malloc(sizeof(int) * res_array_size);
    memset(res, 0, sizeof(int) * res_array_size);
    
    for (int j = 0; j < res_size; j++) // initialize res[] with 1st row of mat
        res[j] = mat[0][j];
    
    for (int i = 1; i < matSize; i++) {
        /* calculate sum from two row */
        int tmp_size = res_size * *matColSize; // size of temp array should be
        int tmp_cnt = 0;
        
        q = (struct pq *)malloc(sizeof(struct pq));
        memset(q, 0, sizeof(struct pq));
        q->heap = (int *)malloc(sizeof(int) * tmp_size);
        memset(q->heap, 0, sizeof(int) * tmp_size);
        
        /* 1. generate temp array from mat[i][] and res[] */
        for (int j = 0; j < *matColSize; j++) // O(m * n * k)
            for (int k = 0; k < res_size; k++)
                q->heap[q->heap_size++] = res[k] + mat[i][j];    
        
        /* 2. make temp array to heap structure with O(N) */
        build_heap(q->heap, q->heap_size); // O(m * k)
        
        /* 3. save res[] with maximum k size, for calulating with next mat[i + 1][] */
        res_size = tmp_size < k ? tmp_size : k; // size of res[] can't be larger than k
        for (int idx = 0; idx < res_size; idx++) // O(m * k * log k)
            res[idx] = heap_pop(q->heap);
        
        free(q->heap);
        free(q);
    }
    return res[k - 1];
}