RandomPool等概率随机快速获取key

等概率随机快速获取key的结构

RandomPool结构

完成以下三个功能:
insert(key):将某个key加入到该结构，做到不重复加入；
delete(key):将原本在结构中的某个key移除；
getRandom(): 等概率随机返回结构中的任何一个key。
【要求】
insert、delete和getRandom方法的时间复杂度都是 $O(1)$。

算法思路

HashSet结构insert、delete操作时间复杂度为 $O(1)$，而现需要等概率随机返回key；
random 等概率生成数值，而HashSet按key索引而不是数值索引，因此想到使用HashMap引入数值；
所以用两个HashMap：key_num_map和num_key_map

• key_num_map维护key值
• num_key_map通过随机生成的数值获取key
  
  

要点：

• count保存key数量

• 对应的num值随着count增加（唯一）；

• key_num_map和num_key_map在移除给定的key时需要调整对应的num和count；
  num为给定key对应的值，fill为count对应的key，——》表示键值对变化后的状态
  key_num_map中:
  key - num ——》移除
  fill - count ——》 fill - num

  num_key_map中:
  num - key ——》num - fill
  count - fill ——》移除
  
  

相应代码

import random

class RandomPool():
    def __init__(self):
        self.key_num_map = {}
        self.num_key_map = {}
        self.count = 0

    def insert(self, key):
        if key not in self.key_num_map:
            self.key_num_map[key] = self.count
            self.num_key_map[self.count] = key
            self.count += 1

    def delete(self, key):
        if key in self.key_num_map:
            self.count -= 1
            num = self.key_num_map[key]  # number the key mapping in key_num_map
            fill = self.num_key_map[self.count]  # last key in num_key_map
            self.key_num_map[fill] = num
            self.num_key_map[num] = fill
            del self.num_key_map[self.count]
            del self.key_num_map[key]

    def get_random(self):
        return self.num_key_map[random.randrange(0, self.count)]

    def __str__(self):
        return ' '.join(self.key_num_map.keys())

# 简单测试
if __name__ == '__main__':
    random_pool = RandomPool()
    random_pool.insert('jiang')
    random_pool.insert('zhi')
    random_pool.insert('qiang')
    print(random_pool)
    print(random_pool.get_random())

    random_pool.insert('jiang')
    print(random_pool)

    random_pool.delete('jiang')
    print(random_pool)
    print(random_pool.get_random())

有任何疑问和建议，欢迎在评论区留言和指正！

感谢您所花费的时间与精力！