FilesExpand file tree

 master

/

randomized_set.py

Copy path

Blame

More file actions

Blame

More file actions

Latest commit

 

History

History

History

70 lines (53 loc) · 1.66 KB

 master

/

randomized_set.py

Copy path

Top

File metadata and controls

• Code
• Blame

70 lines (53 loc) · 1.66 KB

Raw

Copy raw file

Download raw file

Open symbols panel

Edit and raw actions

1

2

3

4

5

6

7

8

9

10

11

12

13

14

15

16

17

18

19

20

21

22

23

24

25

26

27

28

29

30

31

32

33

34

35

36

37

38

39

40

41

42

43

44

45

46

47

48

49

50

51

52

53

54

55

56

57

58

59

60

61

62

63

64

65

66

67

68

69

70

#! /usr/bin/env python3

"""

Design a data structure that supports all following operations

in average O(1) time.

insert(val): Inserts an item val to the set if not already present.

remove(val): Removes an item val from the set if present.

random_element: Returns a random element from current set of elements.

Each element must have the same probability of being returned.

"""

import random

class RandomizedSet():

"""

idea: shoot

"""

def __init__(self):

self.elements = []

self.index_map = {} # element -> index

def insert(self, new_one):

if new_one in self.index_map:

return

self.index_map[new_one] = len(self.elements)

self.elements.append(new_one)

def remove(self, old_one):

if not old_one in self.index_map:

return

index = self.index_map[old_one]

last = self.elements.pop()

self.index_map.pop(old_one)

if index == len(self.elements):

return

self.elements[index] = last

self.index_map[last] = index

def random_element(self):

return random.choice(self.elements)

def __test():

rset = RandomizedSet()

ground_truth = set()

n = 64

for i in range(n):

rset.insert(i)

ground_truth.add(i)

# Remove a half

for i in random.sample(range(n), n // 2):

rset.remove(i)

ground_truth.remove(i)

print(len(ground_truth), len(rset.elements), len(rset.index_map))

for i in ground_truth:

assert(i == rset.elements[rset.index_map[i]])

for i in range(n):

print(rset.random_element(), end=' ')

print()

if __name__ == "__main__":

__test()