Katkat
Back

Multiset with O(1) Insert, Remove, and Count-Weighted Random Pick

SnowflakeSoftware Engineer · Onsite
MediumSoftware EngineerOnsiteSoftware Engineering Fundamentals

Last updated

Design a collection of integers that allows duplicates and supports insert, remove and a random pick weighted by occurrence count, each in average constant time. Tests array and hash map bookkeeping, careful handling of moved elements during removal, and how to test randomized behavior.

Design a data structure that holds a collection of integers in which duplicates are allowed, and that supports each of the following operations in average O(1) time:

  • insert(val) -> bool: add one occurrence of val. Return True if val was not in the collection before this call, and False otherwise.
  • remove(val) -> bool: if val is present, remove one occurrence of it and return True; otherwise return False.
  • get_random() -> int: return a random element of the collection, where the probability of returning a value is proportional to its number of occurrences. If the collection holds [1, 1, 2], get_random returns 1 with probability 2/3 and 2 with probability 1/3.

Write the class, then write a few test cases that exercise it.

Constraints

  • Values fit in a 32-bit signed integer: -2**31 <= val <= 2**31 - 1.
  • At most 2 * 10**5 calls are made in total.
  • get_random is called only when the collection is not empty.

Clarifying Questions

  • Is average (expected) O(1) acceptable, or must every operation be O(1) in the worst case?
  • Should the random source be injectable, so that tests can be deterministic?
  • Must the structure be safe to use from several threads at once?

Hints

Hint 1

Constant-time random pick: Sampling in proportion to occurrence counts is trivial for one kind of container if every occurrence is stored separately. Think about which container that is, and what then makes `remove` hard.

Hint 2

Finding an occurrence: To remove one occurrence of a value quickly, you need to know where at least one of its copies is stored, and to keep that knowledge correct whenever something else moves.

Follow-up questions

  • How would you change the structure so that `get_random` picks each distinct value with equal probability, regardless of its count?
  • How would you test that `get_random` really returns values in proportion to their counts?
  • What changes if several threads call these methods concurrently?
  • How would you add `remove_all(val)`, and what would it cost?