堆实现O(1)时间复杂度查询中位数

最新推荐文章于 2024-08-22 23:39:57 发布

Lhw_666

最新推荐文章于 2024-08-22 23:39:57 发布

阅读量1k

点赞数 2

文章标签：数据结构算法中位数 c++ 堆

本文链接：https://blog.csdn.net/Lhw_666/article/details/113090734

版权

本文介绍了如何使用大根堆和小根堆来动态维护一个序列的中位数，详细阐述了算法流程和关键操作如插入、调整和弹出堆顶元素，并提供了完整的C++代码实现。此算法在插入新数时时间复杂度为O(logN)，查询中位数则为O(1)。

摘要由CSDN通过智能技术生成

思想

准备两个堆，一个大根堆，一个小根堆
对于新来的一个数：
（1）如果大于等于大根堆的堆顶，则放入小根堆
（2）如果小于大根堆的堆顶，则放入大根堆
（3）检查大根堆的大小与小根堆的大小是否相差小于等于1，如果大根堆大小减去小根堆大小大于1，则将大根堆的堆顶弹出放入小根堆；如果小根堆大小减去大根堆大小大于1，则将小根堆的堆顶弹出放入大根堆。即始终保证大根堆保存着较小n/2个数，小根堆保存着较大的n/2个数
求中位数：
（1）如果两个堆的大小相等，则中位数为两个堆的堆顶的平均值
（2）如果两个堆的大小不相等，则中位数为大小较大的堆的堆顶
分析：
每次到来一个新的数，需要O(logN)的时间复杂度实现堆的插入，查询中位数只需O(1)时间复杂度
使用一维数组实现堆的具体思想可以参考这篇文章：
大根堆（数组实现）

代码实现

swap()

用于交换两个下标对应的数

void swap(int arr[], int i, int j){
    int temp = arr[i];
    arr[i] = arr[j];
    arr[j] = temp;
}

heapInsert()

实现堆的插入操作

void heapInsert(int heap[], int newNum, int &heapSize, bool isBigRoot){
    int insertLoc = heapSize;
    heap[insertLoc] = newNum;
    if(isBigRoot){
        while(heap[insertLoc] > heap[(insertLoc - 1) / 2]){
            swap(heap, insertLoc, (insertLoc - 1) / 2);
            insertLoc = (insertLoc - 1) / 2;
        }
    }else{
        while(heap[insertLoc] < heap[(insertLoc - 1) / 2]){
            swap(heap, insertLoc, (insertLoc - 1) / 2);
            insertLoc = (insertLoc - 1) / 2;
        }
    }
    heapSize++;
}

heapify()

当弹出堆顶元素后，实现堆的调整操作

void heapify(int heap[], int loc, int &heapSize, bool isBigHeap){
    int leftChild = 2 * loc + 1;
    if(isBigHeap){
       while(leftChild < heapSize){
            int bigLoc = leftChild + 1 < heapSize && heap[leftChild + 1] > heap[leftChild] ? leftChild + 1 : leftChild;
            if(heap[loc] < heap[bigLoc])
                swap(heap, loc, bigLoc);
            else
                break;
            loc = bigLoc;
            leftChild = 2 * loc + 1;
        }
    }else{
        while(leftChild < heapSize){
            int smallLoc = leftChild + 1 < heapSize && heap[leftChild + 1] < heap[leftChild] ? leftChild + 1 : leftChild;
            if(heap[loc] > heap[smallLoc])
                swap(heap, loc, smallLoc);
            else
                break;
            loc = smallLoc;
            leftChild = 2 * loc + 1;
        }
    }
}

popHeap()

实现弹出堆顶操作，返回堆顶元素

int popHeap(int heap[], int &heapSize, bool isBigHeap){
    int popNum = heap[0];
    swap(heap, 0, --heapSize);
    heapify(heap, 0, heapSize, isBigHeap);
    return popNum;
}

完整代码

#include <iostream>

using namespace std;

void swap(int arr[], int i, int j){
    int temp = arr[i];
    arr[i] = arr[j];
    arr[j] = temp;
}

void heapInsert(int heap[], int newNum, int &heapSize, bool isBigRoot){
    int insertLoc = heapSize;
    heap[insertLoc] = newNum;
    if(isBigRoot){
        while(heap[insertLoc] > heap[(insertLoc - 1) / 2]){
            swap(heap, insertLoc, (insertLoc - 1) / 2);
            insertLoc = (insertLoc - 1) / 2;
        }
    }else{
        while(heap[insertLoc] < heap[(insertLoc - 1) / 2]){
            swap(heap, insertLoc, (insertLoc - 1) / 2);
            insertLoc = (insertLoc - 1) / 2;
        }
    }
    heapSize++;
}

void heapify(int heap[], int loc, int &heapSize, bool isBigHeap){
    int leftChild = 2 * loc + 1;
    if(isBigHeap){
       while(leftChild < heapSize){
            int bigLoc = leftChild + 1 < heapSize && heap[leftChild + 1] > heap[leftChild] ? leftChild + 1 : leftChild;
            if(heap[loc] < heap[bigLoc])
                swap(heap, loc, bigLoc);
            else
                break;
            loc = bigLoc;
            leftChild = 2 * loc + 1;
        }
    }else{
        while(leftChild < heapSize){
            int smallLoc = leftChild + 1 < heapSize && heap[leftChild + 1] < heap[leftChild] ? leftChild + 1 : leftChild;
            if(heap[loc] > heap[smallLoc])
                swap(heap, loc, smallLoc);
            else
                break;
            loc = smallLoc;
            leftChild = 2 * loc + 1;
        }
    }
}

int popHeap(int heap[], int &heapSize, bool isBigHeap){
    int popNum = heap[0];
    swap(heap, 0, --heapSize);
    heapify(heap, 0, heapSize, isBigHeap);
    return popNum;
}

int *bigRootHeap = new int[1000], bigRootHeapSize = 0;
int *smallRootHeap = new int[1000], smallRootHeapSize = 0;

int main()
{
    int n;
    cin >> n;
    for(int i = 0; i < n; i++){
        int newNum;
        cin >> newNum;
        if(newNum < bigRootHeap[0]){
            heapInsert(bigRootHeap, newNum, bigRootHeapSize, true);
        }else{
            heapInsert(smallRootHeap, newNum, smallRootHeapSize, false);
        }
        if(bigRootHeapSize - smallRootHeapSize > 1){
            int popNum = popHeap(bigRootHeap, bigRootHeapSize, true);
            heapInsert(smallRootHeap, popNum, smallRootHeapSize, false);
        }else if(smallRootHeapSize - bigRootHeapSize > 1){
            int popNum = popHeap(smallRootHeap, smallRootHeapSize, false);
            heapInsert(bigRootHeap, popNum, bigRootHeapSize, true);
        }
    }
    float ans = bigRootHeapSize ^ smallRootHeapSize ? bigRootHeapSize > smallRootHeapSize ? bigRootHeap[0] : smallRootHeap[0] : 1.0 * (bigRootHeap[0] + smallRootHeap[0]) / 2;
    cout << ans << endl;
    return 0;
}