Leetcode - 1337. The K Weakest Rows in a Matrix

質問する

2次元配列が与えられ、各行に対応する配列の1つの数がk個の配列を小さい順に出力し、出力される値は行のインデックス番号である.
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)

配列を並べ替えることは可能であるが,O(N log N)よりも性能が優れているわけではない(countig,基数sortなどを除き,最も速い並べ替えアルゴリズムはN logNである).問題はソート後の配列では最大値順がk個しか要求されないため、すべての配列をHeap構造に設定してソートを必要としない場合、O(N)、K回heap pop O(logn)だけでO(k*logn+N)で解決できる.k=Nの最悪の場合はNlognであるが,kが小さいとO(N)に近い.
また,heapify()比較の対象が配列値のみならず,その配列値が同じである場合,異なる優先度で比較する必要があるという問題がしばしば発生する.この場合、比較関数を増やして問題を解決すれば便利になるようです.(アレイ値のサイズを比較する不等式->比較する条件が複雑な場合は、他の比較関数を比較します)
問題では,1の個数(兵士値)が同じであればrowのインデックス値はさらに低い.したがって、以下に示すように、is little()という比較関数が作成されます.この比較関数は、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,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;
}