0
点赞
收藏
分享

微信扫一扫

Python - 构造O(1)时间插入、删除和获取随机元素的集合

爱奔跑的读书者 2022-02-17 阅读 72
python算法

一.引言

设计一个支持在平均 时间复杂度 O(1) 执行以下操作的数据结构。

注意: 允许出现重复元素。

集合包含以下三个功能:

A.insert(val):向集合中插入元素 val。

B.remove(val):当 val 存在时,从集合中移除一个 val。

C.getRandom:从现有集合中随机获取一个元素。每个元素被返回的概率应该与其在集合中的数量呈线性相关。

集合功能示例:

二.实现

集合类型定义为 RanomizedSet,实现了上述三个方法即插入元素,移除元素以及随机获取与元素个数线性相关的元素且时间复杂度为 O(1)。这里保存集合元素采用了 list,判断元素是否存在则使用了 dict 辅助,所以虽然该数据结构上述操作的时间复杂度为 O(1),但是为此空间复杂度也随之增加。

#!/usr/bin/python
# -*- coding: UTF-8 -*-

from random import choice


class RandomizedSet:
    def __init__(self):
        """
        Initialize your data structure here.
        """
        self.dict = {}
        self.list = []

    def insert(self, val: int) -> bool:
        """
        Inserts a value to the set. Returns true if the set did not already contain the specified element.
        """
        self.list.append(val)
        if val in self.dict:
            return False
        self.dict[val] = len(self.list) - 1
        return True

    def remove(self, val: int) -> bool:
        """
        Removes a value from the set. Returns true if the set contained the specified element.
        """
        if val in self.dict:
            last_element, idx = self.list[-1], self.dict[val]
            self.list[idx], self.dict[last_element] = last_element, idx
            self.list.pop()
            del self.dict[val]
            return True
        return False

    def getRandom(self) -> int:
        """
        Get a random element from the set.
        """
        return choice(self.list

三.详解

1.insert(val)

def insert(self, val: int) -> bool:
    """
    Inserts a value to the set. Returns true if the set did not already contain the specified element.
    """
    self.list.append(val)
    if val in self.dict:
        return False
    self.dict[val] = len(self.list) - 1
    return True

向集合中插入元素,由于要求中提到允许出现重复的元素,所以这里开头就执行 list.append 的操作添加元素,这里原题目中给出的代码有问题,无法添加重复元素。List 中添加新元素 val 后,存储 val 在列表中对应的位置,这里通过 dict 保证一个元素只保存一个位置。返回值通过 dict 判断,如果 dict 存在则返回 False,否则返回 True。

2.remove(val)

def remove(self, val: int) -> bool:
    """
    Removes a value from the set. Returns true if the set contained the specified element.
    """
    if val in self.dict:
        last_element, idx = self.list[-1], self.dict[val]
        self.list[idx], self.dict[last_element] = last_element, idx
        self.list.pop()
        del self.dict[val]
        return True
    return False

移除元素中的元素,这里是这个 collection 最复杂的地方(其实也还好),这里的思路类似于冒泡排序一样。

A.获取最后一个元素的值和要删除值的索引

last_element, idx = self.list[-1], self.dict[val]

B.删除值位置的元素替换为列表最后一个值,并修改其在 dict 的索引位置

self.list[idx], self.dict[last_element] = last_element, idx

C.删除该值

上一个步骤修改列表 idx 处的值其实已经将 val 元素删除,这一步主要是更新数据结构,pop 减小列表长度并删除 dict 中多余的索引。

D.返回值

包含元素并移除返回 True,否则返回 False。

3.getRandom(val)

这里 choice 函数来自 random 库,需要 from random import choice 引入,这里如果想自己实现也很简单,只需要随机 [0, len(list)-1] 的数字并从 list 取值即可,随机到每个索引的概率为 1/n,n 为列表长度,所以 p(val) = 1/n * num,num 为该数字在集合中的个数,因此 p(val) 满足概率与其在集合中的数量呈线性相关。

def getRandom(self) -> int:
    """
    Get a random element from the set.
    """
    return choice(self.list)
举报

相关推荐

0 条评论