Leetcode - 1337. The K Weakest Rows in a Matrix

문제

2차원 배열이 주어지고 각 row에 해당하는 배열의 1갯수가 작은 순서대로 k개만큼 출력하기, 출력하는 값은 row의 index번호.
https://leetcode.com/problems/the-k-weakest-rows-in-a-matrix/

Input: mat = 
[[1,1,0,0,0],
 [1,1,1,1,0],
 [1,0,0,0,0],
 [1,1,0,0,0],
 [1,1,1,1,1]], 
k = 3
Output: [2,0,3]
Explanation: 
The number of soldiers in each row is: 
- Row 0: 2 
- Row 1: 4 
- Row 2: 1 
- Row 3: 2 
- Row 4: 5 
The rows ordered from weakest to strongest are [2,0,3,1,4].

해결 - Heap O(N)

배열을 sorting해도 되지만, 그러면 아무리해도 O(N log N)보다 성능이 좋을 수 없다(countig, radix sort등을 제외하면 가장 빠른 정렬알고리즘이 N logN이므로). 문제가 요구하는 것이 sorting된 배열에서 가장큰 값 순서대로 딱 k개만 요구하므로, 모든 배열을 정렬할 필요없이 Heap 구조로 만든 뒤O(N), k번만 heap pop O(logN) 한다면 O(k*logN + N)으로 해결 가능하다. k == N인 최악의 경우 NlogN 이 되겠지만 k가 작은경우 O(N)에 더 가까울 것이다.

추가로 종종 heapify()에서 비교하는 대상이 단순히 배열값뿐만 아니라, 그 배열값이 같은경우는 다른 우선순위로 비교하라는 문제를 종종 만나는데 많이 햇갈릴 수 있다. 그런 경우는 문제를 비교함수를 추가로 만들어 해결하면 편한것 같다. (배열값 크기의 부등호 비교 -> 비교해야할 조건이 보다 복잡하다면 추가 비교함수로 비교)

문제에서는 1의 갯수(soldiers값)가 동일하다면 row의 인덱스가 더 낮은값이 더 작은값이라고 명시되었다. 그래서 아래와 같이 is_smaller()라는 비교함수를 만들었다. heap이 Min Heap인 경우 heapify에서 가장 작은 값을 찾을때 이 비교함수를 사용하면 된다.

bool is_smaller(struct pq a, struct pq b)
{
    if (a.soldiers < b.soldiers || (a.soldiers == b.soldiers && a.row < b.row))
        return true;
    return false;
}

void heapify(struct pq *q, int p_idx, int size)
{ ...중략
    if (l_idx < size && is_smaller(q[l_idx], q[min_idx]))
        min_idx = l_idx;

만약 다른문제에서도 비교 조건이 배열값의 부등호 크기만이라고 해도 아래와같이 compare()함수로 구현한다고 생각하면 앞으로 더 복잡한 조건의 비교를 필요로하는 문제를 만나도 쉽게 해결 가능할것같다. 더 범용적으로는 compare() 함수를 heapify()의 인자로 함수포인터로 받는것도 좋을듯.

bool compare(int a, int b)
{
	return (a < b) ? true: false;
}

void heapify(int arr[], int p_idx, int size)
{ ...중략
    if (l_idx < size && compare(arr[l_idx], arr[min_idx]))
        min_idx = l_idx;

전체 코드는 다음과 같다. 결과는 Runtime: 12 ms, faster than 97.84%. Memory Usage: 7 MB, less than 21.55% of of C online submissions.

struct pq {
    int soldiers;
    int row;
};
struct pq *q;
int heap_size;

void swap(struct pq *q, int a, int b)
{
    struct pq temp = q[a];
    q[a] = q[b];
    q[b] = temp;
}

bool is_smaller(struct pq a, struct pq b)
{
    if (a.soldiers < b.soldiers || (a.soldiers == b.soldiers && a.row < b.row))
        return true;
    return false;
}

void heapify(struct pq *q, int p_idx, int size)
{
    int min_idx = p_idx;
    int l_idx = p_idx * 2 + 1;
    int r_idx = p_idx * 2 + 2;
    
    // find min, and if q[i].soldiers are same, then compare q[i].row 
    if (l_idx < size && is_smaller(q[l_idx], q[min_idx]))
        min_idx = l_idx;
    if (r_idx < size && is_smaller(q[r_idx], q[min_idx]))
        min_idx = r_idx;

    if (min_idx != p_idx) { /* base case */
        swap(q, min_idx, p_idx);
        heapify(q, min_idx, size);
    }
}

void build_heap(struct pq *q, int size)
{
    for (int i = size / 2 -1; i >= 0; i--)
        heapify(q, i, size);
}

int heap_pop(struct pq *q)
{
    int ret = q[0].row;
    heap_size--;
    swap(q, 0, heap_size);
    heapify(q, 0, heap_size);
    return ret;
}

void print_arr(struct pq *q, int size)
{
    for (int i = 0; i < size; i++)
        printf("[%d]%d ", q[i].row, q[i].soldiers);
    printf("\n");
}

/**
 * Note: The returned array must be malloced, assume caller calls free().
 */
int* kWeakestRows(int** mat, int matSize, int* matColSize, int k, int* returnSize){
    *returnSize = k;
    heap_size = matSize;
    int * ret = (int *)malloc(sizeof(int) * k);
    q = (struct pq *)malloc(sizeof(struct pq) * 101);
    memset(q, 0, sizeof(struct pq) * 101);
    for (int i = 0; i < matSize; i++) {
        int cnt = 0;
        while (cnt < *matColSize && mat[i][cnt] != 0)
            cnt++;
        q[i].row = i; // {0, 1, 2, 3, 4, 5} idx 
        q[i].soldiers = cnt;
    }
    build_heap(q, heap_size);
    for (int i = 0; i < k; i++)
        ret[i] = heap_pop(q);
        
    return ret;
}