当前位置：首页 > article >正文

面试经典问题 —— 最大/小前K个数问题（top - K）问题

article 2024/12/26 22:20:42

常见思路

要选出最小的前K个数首先我们会想到排排升序建大堆，排降序建小堆

一个直观的想法是使用（小根堆），起始将所有元素放入堆中，然后再从堆中
 取出k 个元素并「顺序」构造答案。

你写出了用小堆解决的代码


void swap(int* a, int* b) {
    int t = *a;
    *a = *b;
    *b = t;
}


void minHeapify(int arr[], int n, int i) {
    int smallest = i;
    int left = 2 * i + 1;
    int right = 2 * i + 2; 


    if (left < n && arr[left] < arr[smallest])
        smallest = left;

   
    if (right < n && arr[right] < arr[smallest])
        smallest = right;


    if (smallest != i) {
        swap(&arr[i], &arr[smallest]);
        
        minHeapify(arr, n, smallest);
    }
}


void buildMinHeap(int arr[], int n) {
    // Index of last non-leaf node
    int startIdx = (n / 2) - 1;

    for (int i = startIdx; i >= 0; i--)
        minHeapify(arr, n, i);
}


int extractMin(int arr[], int* n) {
    if (*n <= 0)
        return INT_MAX;
    if (*n == 1) {
        (*n)--;
        return arr[0];
    }


    int root = arr[0];
    arr[0] = arr[*n - 1];
    (*n)--;
    minHeapify(arr, *n, 0);

    return root;
}

void findKSmallest(int arr[], int n, int k) {

    buildMinHeap(arr, n);

    printf("The k smallest elements are: ");
    for (int i = 0; i < k; i++) {
        printf("%d ", extractMin(arr, &n));
    }
    printf("\n");
}

分析一下不难发现
建堆的复杂度是 O（N）
每次提取的复杂度是 O(logN)
总的复杂度O（NlogN）

面试官看到你的代码后，并不满意并要求你进行优化

更优的解法（面试官喜欢的）

使用（大根堆）!
当处理到原始 arr[i] 时，根据堆内元素个数，以及其与堆顶元素的关系分情况讨论：
堆内元素不足 k 个：直接将 arr[i] 放入堆内；
堆内元素为 k 个：根据 arr[i] 与堆顶元素的大小关系分情况讨论：
arr[i]>=heapTop：arr[i] 不可能属于第 k 小数（已有 k 个元素在堆中），直接丢弃 arr[i]；
arr[i]<heapTop：arr[i] 可能属于第 k 小数，弹出堆顶元素，并放入 arr[i]。

在这里插入图片描述

你出写了如下代码。

void swap(int* a, int* b) {
    int t = *a;
    *a = *b;
    *b = t;
}


void maxHeapify(int arr[], int n, int i) {
    int largest = i;
    int left = 2 * i + 1;
    int right = 2 * i + 2; 

    // 如果左子节点大于当前节点，则更新最大值索引
    if (left < n && arr[left] > arr[largest])
        largest = left;

    // 如果右子节点大于当前节点，则更新最大值索引
    if (right < n && arr[right] > arr[largest])
        largest = right;

    // 如果最大值索引不是当前节点，则交换并递归调整子树
    if (largest != i) {
        swap(&arr[i], &arr[largest]);
        maxHeapify(arr, n, largest);
    }
}


void buildMaxHeap(int arr[], int n) {
    // 最后一个非叶子节点的索引
    int startIdx = (n / 2) - 1;

    // 从最后一个非叶子节点开始，逆序遍历并调整每个节点
    for (int i = startIdx; i >= 0; i--)
        maxHeapify(arr, n, i);
}

// 提取堆顶元素并调整堆的函数
int extractMax(int arr[], int* n) {
    if (*n <= 0)
        return INT_MAX; // 堆为空时返回最大整数值
    if (*n == 1) {
        (*n)--;
        return arr[0]; // 堆只有一个元素时直接返回
    }

    // 存储并移除堆顶元素
    int root = arr[0];
    arr[0] = arr[*n - 1];
    (*n)--;
    maxHeapify(arr, *n, 0); // 调整堆以保持大根堆性质

    return root;
}


void findKSmallest(int arr[], int n, int k) {
    // 构建包含前k个元素的大根堆
    int heap[k];
    for (int i = 0; i < k; i++) {
        heap[i] = arr[i];
    }
    buildMaxHeap(heap, k);

    // 遍历剩余元素
    for (int i = k; i < n; i++) {
        // 如果当前元素小于堆顶元素，则替换堆顶元素并重新调整堆
        if (arr[i] < heap[0]) {
            heap[0] = arr[i];
            maxHeapify(heap, k, 0);
        }
    }
}