TopK问题

TopK问题

syndi

2019 年 05 月 24 日

988 次浏览

暂无评论

1219字数

1. 问题描述

从一个规模巨大的数据集中找到最大/最小的K个元素。

2. 直接排序

最粗暴的方法，直接排序后，想要最大或者最小的K个元素直接取就行了。我使用了快速排序，平均情况的时间复杂度是O(NlogN)，最坏情况O(N^2)。常规思路没什么新意，并且需要能把全部数据集一次性加载到内存中。

void Swap(int* x, int* y) {
    int tmp = *x;
    *x = *y;
    *y = tmp;
}

int partitionD(int* arr, int left, int right) {
    int pivot = arr[left];  // base
    int l = left;
    int r = right;

    while (l < r) {
        while (l < r && arr[r] >= pivot) r--;
        while (l < r && arr[l] <= pivot) l++;
        if (l < r) {
            Swap(&arr[l], &arr[r]);
        }
    }
    Swap(&arr[left], &arr[l]);
    return l;
}

int partition(int* arr, int left, int right) {
    int pivot = arr[left];  // base
    int l = left + 1;
    int r = left + 1;

    while (r <= right) {
        if (arr[r] < pivot) {
            Swap(&arr[r], &arr[l++]);
        }
        r++;
    }
    Swap(&arr[left], &arr[l - 1]);
    return l - 1;
}

void __quick_sort(int* arr, int left, int right) {
    if (left < right) {
        int pivotIndex = partitionD(arr, left, right);
        __quick_sort(arr, left, pivotIndex - 1);
        __quick_sort(arr, pivotIndex + 1, right);
    }
}

void quick_sort(int* arr, int size) {
    __quick_sort(arr, 0, size - 1);
}

int* TopKSort(int* arr, int size, int k) {
    quick_sort(arr, size);

    int* res = (int *)malloc(sizeof(int) * k);
    for (int i = 0; i < k; i++) {
        res[i] = arr[i];
    }

    return res;
}

3. 维护一个大小为K的数据集

维护一个大小为K的数据集。假如要找K个最小的元素，则先把前K个元素加入数据集中。然后往后遍历剩余元素，每处理一个新元素就和数据集中的最大元素比较，如果比最大元素还小，那么最大元素一定不可能是K个最小元素之一，于是用新元素替换掉最大元素，重复上述步骤即可。这个方法时间复杂度为O(KN)，而且不需要一次性加载全部数据到内存。

// 找到最大元素下标
int findBiggest(int* arr, int size) {
    int res = 0;
    for (int i = 1; i < size; i++) {
        if (arr[i] > arr[res]) {
            res = i;
        }
    }

    return res;
}

int* TopKArray(int* arr, int size, int k) {
    int* res = (int *)malloc(sizeof(int) * k);
    memcpy(res, arr, sizeof(int) * k);

    // 遍历后续元素
    for (int index = k; index < size; index++) {
        int biggest = findBiggest(res, k);
        if (arr[index] < res[biggest]) {
            res[biggest] = arr[index];
        }
    }

    return res;
}

4. 最大堆/最小堆

假如需要找K个最小的数。那么可以用前K个数建一个最大堆，可知此时堆顶元素为堆中最大元素，则处理后续元素的时候只需要和堆顶元素比较，假如比堆顶元素小，则可知堆顶元素一定不是最小的K个元素之一，于是用这个新数据替换掉堆顶元素，然后向下调整堆顶元素再建堆重复上述过程。

这个方法时间复杂度为O(NlogK)，并且不需要一次性把全部元素加载到内存中。

void Swap(int* x, int* y) {
    int tmp = *x;
    *x = *y;
    *y = tmp;
}

void AdjustDown(int* heap, int size, int root) {
    int parent = root;
    int child = parent * 2 + 1;

    while (child < size) {
        // 找出最小孩子节点
        if (child + 1 < size && heap[child + 1] > heap[child]) {
            child++;
        }

        if (heap[child] > heap[parent]) {
            Swap(&heap[child], &heap[parent]);
            parent = child;
            child = parent * 2 + 1;
        }
        else {
            break;
        }
    }
}

int* HeapCreate(int* arr, int k) {
    int* heap = (int *)malloc(sizeof(int) * k);
    memcpy(heap, arr, sizeof(int) * k);

    // 建堆
    int root = (k - 2) / 2;  // 第一个需要调整的子树
    for (; root >= 0; root--) {
        AdjustDown(heap, k, root);
    }

    return heap;
}

// K个最小的数建大堆，K个最大的数建小堆
int* TopKHeap(int* arr, int size, int k) {
    // 用数组前K个数建个最大堆，并返回
    int* heap = HeapCreate(arr, k);

    // 遍历后续数据
    for (int index = k; index < size; index++) {
        if (arr[index] < heap[0]) {  // 如果比堆顶元素小
            heap[0] = arr[index];
            AdjustDown(heap, k, 0);
        }
    }

    return heap;
}

TopK问题