数据结构算法刷题（15）高效、等概率地随机获取元素，数据流的中位数

最新推荐文章于 2026-06-20 21:44:23 发布

原创最新推荐文章于 2026-06-20 21:44:23 发布 · 179 阅读

0 ·

本内容遵循CC 4.0 BY-SA版权协议

标签

#数据结构

如果想「等概率」且「在 O(1) 的时间」取出元素，一定要满足：底层用数组实现，且数组必须是紧凑的。

但如果用数组存储元素的话，插入，删除的时间复杂度怎么可能是 O(1) 呢？

对数组尾部进行插入和删除操作不会涉及数据搬移，时间复杂度是 O(1)。

O(1) 的时间删除数组中的某一个元素val，可以先把这个元素交换到数组的尾部，然后再pop掉。

class RandomizedSet:

def __init__(self):

import collections

import random

self.indexdict = collections.defaultdict() #用来存数组的下标

self.setlist = [] #用来存数据

def insert(self, val: int) -> bool:

if val not in self.setlist:

self.setlist.append(val) #把元素插到末尾

self.indexdict[val] = len(self.setlist) - 1 #记录元素的下标

return True

else:

return False

def remove(self, val: int) -> bool:

if val in self.setlist:

index_key = self.indexdict[val] #找到该元素对应的下标

a = self.setlist[-1]

self.indexdict[a] = index_key #将交换数据的下标更新成当前元素下标

self.setlist[index_key] = a #把数组最后一位放到index_key的地方

self.setlist.pop() #删除最后一位

self.indexdict.pop(val) #删除对应元素

return True

else:

return False

def getRandom(self) -> int:

a = int(random.random()*len(self.setlist)) #随机取一个下标

return self.setlist[a]

整个取随机数的思想就是哈希映射：将要取的范围想象成一个数组，长度是n，有m个数是黑名单不能取，所以能取的数的范围就是n-m，以n-m为分界线，做一个哈希映射（字典），遍历黑名单，如果该元素在分界线内：if black < n-m，就取第一个不在黑名单的分界线作为映射：

while i in blacklist: i+=1,map(black) = i

每次交换之后，i向后移动。那些不在该分界线范围内的黑名单不需要考虑。

最后在map中取元素，如果随机到的值在map的key中，就取其映射到的值，否则就返回这个随机值。

class Solution:

def __init__(self, n: int, blacklist: List[int]):

m = len(blacklist) #取不到的长度

self.size = n - m #可以取的长度

self.map = dict() #存放映射的字典，key是黑名单中的数字，value是被映射到的数

i = self.size

bset = set(blacklist) #python中set运行比list快

for black in blacklist:

if black < self.size: #那些不在该分界线范围内的黑名单不需要考虑

while i in bset:

i += 1 #取第一个不在黑名单中的i值

self.map[black] = i #映射

i += 1 #映射之后，i向前移动

def pick(self) -> int:

return self.map.get(i:=random.randint(0,self.size-1),i)

我们必然需要有序数据结构，本题的核心思路是使用两个优先级队列。

中位数是有序数组最中间的元素算出来的对吧，我们可以把「有序数组」抽象成一个倒三角形，宽度可以视为元素的大小，那么这个倒三角的中部就是计算中位数的元素对吧：

不仅要维护large和small的元素个数之差不超过 1，还要维护large堆的堆顶元素要大于等于small堆的堆顶元素。

想要往large里添加元素，不能直接添加，而是要先往small里添加，然后再把small的堆顶元素加到large中；向small中添加元素同理。

最后找中位数时，看两个堆的长度，如果一样长，取出堆顶相加除以2，否则取较长的堆顶。

如下代码：首先我实现了一个大根堆和一个小根堆：

class MaxTree: #大根堆
    def __init__(self):
        self.tree1 = ['*'] #数组的第一个是空着不用的
    def getFather(self,index):
        return int(index/2)
    def getLeft(self,index):
        return index*2
    def getRight(self,index):
        return index*2+1
    def less(self,a,b):
        return 1 if a>b else 0
    def swim(self,k): #上浮第k个元素
        while k > 1 and self.less(self.tree1[k],self.tree1[self.getFather(k)]):
            change = self.tree1[self.getFather(k)]
            self.tree1[self.getFather(k)] = self.tree1[k]
            self.tree1[k] = change
            k = self.getFather(k)
    def sink(self,k):
        while self.getLeft(k) < len(self.tree1):
            bigger = self.getLeft(k) #假设左孩子大
            if self.getRight(k) < len(self.tree1) and self.less(self.tree1[self.getRight(k)],self.tree1[self.getLeft(k)]):# 右孩子存在而且右孩子大
                bigger = self.getRight(k)
            if self.tree1[k] < self.tree1[bigger]:
                change = self.tree1[k]
                self.tree1[k] = self.tree1[bigger]
                self.tree1[bigger] = change
                k = bigger
            else:
                break
    def add(self,a):
        self.tree1.append(a)
        self.swim(len(self.tree1)-1)
    def delmax(self):
        self.tree1[1] = self.tree1[len(self.tree1)-1]
        self.tree1.pop(len(self.tree1)-1)
        self.sink(1)

class MinTree: #小根堆
    def __init__(self):
        self.tree1 = ['*'] #数组的第一个是空着不用的
    def getFather(self,index):
        return int(index/2)
    def getLeft(self,index):
        return index*2
    def getRight(self,index):
        return index*2+1
    def less(self,a,b):
        return 1 if a > b else 0
    def swim(self,k): #上浮第k个元素
        while k > 1 and self.less(self.tree1[self.getFather(k)],self.tree1[k]):
            change = self.tree1[self.getFather(k)]
            self.tree1[self.getFather(k)] = self.tree1[k]
            self.tree1[k] = change
            k = self.getFather(k)
    def sink(self,k):
        while self.getLeft(k) < len(self.tree1):
            smaller = self.getLeft(k) #假设左孩子小
            if self.getRight(k) < len(self.tree1) and self.less(self.tree1[self.getLeft(k)],self.tree1[self.getRight(k)]):
                smaller = self.getRight(k)
            if self.less(self.tree1[k],self.tree1[smaller]):
                change = self.tree1[k]
                self.tree1[k] = self.tree1[smaller]
                self.tree1[smaller] = change
                k = smaller
            else:
                break
    def add(self,a):
        self.tree1.append(a)
        self.swim(len(self.tree1)-1)
    def delmin(self):
        self.tree1[1] = self.tree1[len(self.tree1)-1]
        self.tree1.pop(len(self.tree1)-1)
        self.sink(1)

class MedianFinder:

    def __init__(self):
        self.queue_large = MaxTree() #存放较小的数值，堆顶元素最大
        self.queue_small = MinTree() # 存放较大的数值，堆顶元素最小

    def addNum(self, num: int) -> None:
        if len(self.queue_large.tree1) <= len(self.queue_small.tree1):#应该large加
            self.queue_small.add(num) 
            self.queue_large.add(self.queue_small.tree1[1])
            self.queue_small.delmin()
        else:
            self.queue_large.add(num)
            self.queue_small.add(self.queue_large.tree1[1])
            self.queue_large.delmax()

    def findMedian(self) -> float:
        if len(self.queue_large.tree1) == len(self.queue_small.tree1):
            return (self.queue_large.tree1[1] + self.queue_small.tree1[1])/2
        elif len(self.queue_large.tree1) > len(self.queue_small.tree1):
            return self.queue_large.tree1[1]
        else:
            return self.queue_small.tree1[1]