Leetcode - 2099. Find Subsequence of Length K With the Largest Sum

문제

주어진 배열에서 k 개의 subsequece 중에서 그 합이 가장 큰 조합을 출력하라. 기존 배열의 순서가 바뀌면 안됨.

Input: nums = [-1,-2,3,4], k = 3
Output: [-1,3,4]
Explanation: 
The subsequence has the largest sum of -1 + 3 + 4 = 6.

https://leetcode.com/problems/find-subsequence-of-length-k-with-the-largest-sum/

해결 O(K*N)

배열에서 가장 큰수 k개 출력하는건 시간복잡도 O(N) + O(k logN)으로 해결되는 전형적인 heap 문제인데, 이 문제는 기존의 인덱스 순서가 바뀌면 안된다는 조건이 있다. 이부분 때문에 O(KN) 이 되었는데 최악의경우 O(N^2)이 될수 있다.

일단 지난번에 풀었던 해결방식대로 min heap을 k개 남기는 방식으로 풀었다. 남은 k개의 heap배열이 합이 가장 큰 subsequence가 된다. 하지만 순서가 주어진 배열의 index 순서가 아니기 때문에 이를 맞추기 위해서 정렬을 했고 O(K*N)이 되었다. add_num 이후 정렬하는 부분을 아래와 같이 했는데 최악의 경우 O(N^2)이 될수 있다. 아주 마음에 안드는 코드.

    for (int i = 0; i < numsSize; i++) // O(N * logN)
        add_num(q->heap, q->heap_size, nums[i], i);
        
    for (int n = 0; n < numsSize; n++) {
        for (int i = 0; i < k; i++) {
            int cidx = q->heap[i].idx;
            if (n == cidx)
                ret[retcnt++] = nums[cidx]; // FIXME: idx 순서
        }
    }
    return ret;

해결 O(K logN)

위의 정렬하는 부분을 개선. k만큼 heap을 구축 한 뒤에, 다시 idx값을 기준으로 min heapify를 했고. pop한값을 순차적으로 ret[] 에 저장했다. heapify O(N), pop O(k logN)으로 해결됨. 최악의경우 O(NlogN)

heapify_down 매개변수로 함수포인터를 전달한것도 하나의 관전 포인트.

    for (int i = 0; i < numsSize; i++) // O(N * logN)
        add_num(q->heap, q->heap_size, nums[i], i);
    
    build_heap(q->heap, q->heap_size, compare_idx); // O(N)
    for (int i = 0; i < k; i++) // O(K * logN)
        ret[i] = heap_pop(q->heap, compare_idx).val;
    return ret;

제출 결과는 다음과 같음.
Runtime: 7 ms, faster than 93.10%
Memory Usage: 7.1 MB, less than 48.28%

전체 소스코드

struct elem {
    int idx;
    int val;
};

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

void swap(struct elem arr[], int a, int b)
{
    struct elem temp = arr[a];
    arr[a] = arr[b];
    arr[b] = temp;
}

void heapify_up(struct elem arr[], int curr_idx)
{
    int p_idx = (curr_idx - 1) / 2;
    
    if (curr_idx < 1)
        return;
    if (arr[curr_idx].val < arr[p_idx].val) {
        swap(arr, curr_idx, p_idx);
        heapify_up(arr, p_idx);
    }
}

void heap_push(struct elem arr[], int val, int idx)
{
    arr[q->heap_size].val = val;
    arr[q->heap_size].idx = idx;
    q->heap_size++;
    heapify_up(arr, q->heap_size - 1);
}

bool compare_val(struct elem a, struct elem b)
{
    if (a.val < b.val)
        return true;
    else
        return false;
}

bool compare_idx(struct elem a, struct elem b)
{
    if (a.idx < b.idx)
        return true;
    else
        return false;
}

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

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

struct elem heap_pop(struct elem arr[], bool (*cmp)(struct elem, struct elem))
{
    struct elem ret = arr[0];
    swap(arr, 0, q->heap_size -1);
    q->heap_size--;
    heapify_down(arr, 0, q->heap_size, cmp);
    return ret;
}

void add_num(struct elem arr[], int size, int val, int idx)
{
    if (size < g_k) {
        heap_push(arr, val, idx);
    } else if (val > arr[0].val) {
        arr[0].val = val;
        arr[0].idx = idx;
        heapify_down(arr, 0, size, compare_val);
    }
}

// max heap -> heapify - heapify O(N) / pop O(K * logN) 
// min heap with k size - insert O(N*logN) / pop O(K)
int* maxSubsequence(int* nums, int numsSize, int k, int* returnSize){
    q = (struct pq *)malloc(sizeof(struct pq));
    memset(q, 0, sizeof(struct pq));
    q->heap = (struct elem *)malloc(sizeof(struct elem) * k);
    memset(q->heap, 0, sizeof(struct elem) * k);
    *returnSize = k;
    int *ret = (int *)malloc(sizeof(int) * k);
    g_k = k;
    
    for (int i = 0; i < numsSize; i++) // O(N * logN)
        add_num(q->heap, q->heap_size, nums[i], i);
    
    build_heap(q->heap, q->heap_size, compare_idx); // O(N)
    for (int i = 0; i < k; i++) // O(K * logN)
        ret[i] = heap_pop(q->heap, compare_idx).val;
    return ret;
}