profile photo

topK问题的解决方案

2023-09-22 23:32:23

topK问题

最基础的topK问题可以描述如下： 给你一组数字，让你选出第 K 大的数字。

解决方案（not for real world）

1.使用快速排序的选择方法

int quickselect(vector<int> &nums, int l, int r, int k) {
        if (l == r)
            return nums[k];
        int partition = nums[l], i = l - 1, j = r + 1;
        while (i < j) {
            do i++; while (nums[i] < partition);
            do j--; while (nums[j] > partition);
            if (i < j)
                swap(nums[i], nums[j]);
        }
        if (k <= j)return quickselect(nums, l, j, k);
        else return quickselect(nums, j + 1, r, k);
    }

2.BFPRT(Median of medians) 算法

只要系统了解过算法的人应该都知道，对于快速排序，在某些 case 下其时间复杂度将会退化到 O(n^2) 级别。为了解决这种情况，一个比较好的方法就是随机化选择主元。不过我们可以想一下快排的过程，如果我们每次选择的主元都是中位数就好了，这样两边是平衡的，快排就不会退化。但是选择中位数带来的时间开销也要考虑，所以如何更快的选择出中位数便成了一个重要的问题。BFPRT 算法会寻找一个近似中位数，这样虽然无法达到最佳效果但减少了寻找该中位数的时间开销。 该算法会将所有数据划分为 5 个一组（假设总数据数为 N ）。这样一来，我们会得到 N / 5 组数据，找出每组数据的中位数（可以使用插入排序）紧接着递归调用 BFPRT 算法计算这些数据的中位数。（这就是 median of medians 的由来）。紧接着我们利用选出的中位数的中位数作为主元，递归调用BFPRT直到找到 topK 的数。

int bfprt(std::vector<int>& arr, int l, int r, int i) {
    if (l == r) {
        return arr[l];
    }
    int pivot = medianOfMedians(arr, l, r);
    std::vector<int> pivotRange = partition(arr, l, r, pivot);
    if (l + i >= pivotRange[0] && p + i <= pivotRange[1]) {
        return arr[pivotRange[0]];
    } else if (l + i < pivotRange[0]) {
        return bfprt(arr, l, pivotRange[0] - 1, i);
    } else {
        return bfprt(arr, pivotRange[1] + 1, r, i + l - pivotRange[1] - 1);
    }
}

int medianOfMedians(std::vector<int>& arr, int l, int r) {
    int num = r - l + 1;
    int extra = num % 5 == 0 ? 0 : 1;
    std::vector<int> medians(num / 5 + extra);
    for (int i = 0; i < medians.size(); i++) {
        medians[i] = computeMedian(arr, l + i * 5, std::min(l + i * 5 + 4, r));
    }
    return bfprt(medians, 0, medians.size() - 1, medians.size() / 2);
}

std::vector<int> partition(std::vector<int>& arr, int p, int r, int pivot) {
    int small = l - 1;
    int equal = 0;
    int temp;
    for (int j = l; j <= r; j++) {
        if (arr[j] < pivot) {
            small++;
            temp = arr[small];
            arr[small] = arr[j];
            if (equal > 0) {
                arr[j] = arr[small + equal];
                arr[small + equal] = temp;
            } else {
                arr[j] = temp;
            }
        } else if (arr[j] == pivot) {
            equal++;
            temp = arr[j];
            arr[j] = arr[small + equal];
            arr[small + equal] = temp;
        }
    }
    return {small + 1, small + equal};
}

int computeMedian(std::vector<int>& arr, int begin, int end) {
    std::sort(arr.begin() + begin, arr.begin() + end + 1);
    return arr[begin + (end - begin) / 2];
}

3.利用堆排序

我们可以建立一个大根堆，把所有元素插入后，删除k-1个元素，堆顶便是所需元素。(这个方法还有一个优势，可以在线操作)

void maxHeapify(vector<int>& a, int i, int heapSize) {
        int l = i * 2 + 1, r = i * 2 + 2, largest = i;
        if (l < heapSize && a[l] > a[largest]) {
            largest = l;
        } 
        if (r < heapSize && a[r] > a[largest]) {
            largest = r;
        }
        if (largest != i) {
            swap(a[i], a[largest]);
            maxHeapify(a, largest, heapSize);
        }
    }

    void buildMaxHeap(vector<int>& a, int heapSize) {
        for (int i = heapSize / 2; i >= 0; --i) {
            maxHeapify(a, i, heapSize);
        } 
    }

    int findKthLargest(vector<int>& nums, int k) {
        int heapSize = nums.size();
        buildMaxHeap(nums, heapSize);
        for (int i = nums.size() - 1; i >= nums.size() - k + 1; --i) {
            swap(nums[0], nums[i]);
            --heapSize;
            maxHeapify(nums, 0, heapSize);
        }
        return nums[0];
    }

真实世界中的topk

对于真实世界中的topk，我们可能要考虑更多的问题。例如说内存是否可以放下所有的数据；数据是一次性全部给出还是以流的形式给出；数据是否有其他特性。所以当我们面对真实世界的情况时，最好不要直接随意套用这些“官解”，最重要的还是理解思路。

refrence

BFPRT——Top k问题的终极解法 wikipedia Median of medians Thmos.H.Cormen et al. Introduction to Algorithms p.213~222