技術標籤：C++c++

問題

在刷題或者面試的過程中，我們經常會遇到這樣一種題目，給一個數組，求最大或者最小的 k 個數，或者第 k 大或者第 k 小的數字

解法

看到這樣一個題目，首先想到的就是大小堆，建立大小為 k 的大根堆或者小根堆，然後不斷的更新堆，下面以求最大的 k 個數字為例，具體程式碼如下

void BuildHeap(std::vector<int>& arr, int index, int len) {
    int i = index * 2 + 1;
    while (i < len) {
        if ((i + 1) < len && arr[i + 1] < arr[i]) {
            i = i + 1;
        }
        if (arr[i] > arr[index]) {
            break;
        }
        std::swap(arr[i], arr[index]);
        index = i;
        i = 2 * index + 1;
    }
}

void MakeHeap(std::vector<int>& arr) {
    int mid = arr.size() / 2;
    for (int i = mid; i >= 0; --i) {
        BuildHeap(arr, i, arr.size());
    }
}

std::vector<int> MaxK(int* arr, int len, int k) {
    std::vector<int> vec;
    for (int i = 0; i < k; ++i) {
        vec.push_back(arr[i]);
    }
    MakeHeap(vec);
    for (int i = k; i < len; ++i) {
        if (arr[i] > vec[0]) {
            vec[0] = arr[i];
            BuildHeap(vec, 0, k);
        }
    }
    return vec;
}
int main(void) {
    int arr[] = {1,8,3,4,5};
    std::vector<int> val = MaxK(arr, sizeof(arr) / sizeof(arr[0]), 3);
    for (int i = 0; i < 3; ++i) {
        std::cout << val[i] << " ";
    }
    std::cout << std::endl;
    getchar();
    return 0;
}
 

另一種解法 std::priority_queue

C++ STL 有一個優先順序佇列 std::priority_queue，它的內部就是使用堆來實現的，接下來我們使用 std::priority_queue 在來一次

std::vector<int> MaxK(int* arr, int len, int k) {
    std::priority_queue<int, std::vector<int> , std::greater<int>> q{arr, arr+k};
    for (int i = k; i < len; ++i) {
        if (arr[i] > q.top()) {
            q.push(arr[i]);
            q.pop();
        }
    }
    std::vector<int> result;
    while(q.size() > 0) {
        result.push_back(q.top());
        q.pop();
    }
    return result;
}
int main(void) {
    int arr[] = {1,8,3,4,5};
    std::vector<int> val = MaxK(arr, sizeof(arr) / sizeof(arr[0]), 3);
    for (int i = 0; i < val.size(); ++i) {
        std::cout << val[i] << " ";
    }
    std::cout << std::endl;
    getchar();
    return 0;
}
 

程式碼瞬間精簡了很多，可見優先順序佇列還是很好使用的

關於優先順序佇列，可以參考這一篇文章：http://c.biancheng.net/view/480.html

怎麼建立優先順序佇列

1、預設的建立

std::priority_queue<int> q{ arr, arr + k };

上面程式碼表示建立了一個優先順序佇列，初始化元素為陣列的前 k 個元素，然後裡面的元素從大到小排列，如果想要計算最小的 k 個數，或者第 k 小的陣列，就需要按照這種方式建立堆

2、降序排列的優先順序佇列

std::priority_queue<int, std::vector<int> , std::greater<int>> q{arr, arr+k};

表示我們建立一個優先順序佇列，初始化元素為陣列的前 k 個元素，排序的方式為從小到大排序，適用於求前 k 大的數字，或者第 k 大的數字