LeetCode in Kotlin

381. Insert Delete GetRandom O(1) - Duplicates allowed

Hard

RandomizedCollection is a data structure that contains a collection of numbers, possibly duplicates (i.e., a multiset). It should support inserting and removing specific elements and also removing a random element.

Implement the RandomizedCollection class:

You must implement the functions of the class such that each function works on average O(1) time complexity.

Note: The test cases are generated such that getRandom will only be called if there is at least one item in the RandomizedCollection.

Example 1:

Input

["RandomizedCollection", "insert", "insert", "insert", "getRandom", "remove", "getRandom"]
[[], [1], [1], [2], [], [1], []]

Output: [null, true, false, true, 2, true, 1]

Explanation:

RandomizedCollection randomizedCollection = new RandomizedCollection(); 
randomizedCollection.insert(1);     // return true since the collection does not contain 1. 
                                    // Inserts 1 into the collection. 
randomizedCollection.insert(1);     // return false since the collection contains 1. 
                                    // Inserts another 1 into the collection. Collection now contains [1,1]. 
randomizedCollection.insert(2);     // return true since the collection does not contain 2. 
                                    // Inserts 2 into the collection. Collection now contains [1,1,2]. 
randomizedCollection.getRandom();   // getRandom should: 
                                    // - return 1 with probability 2/3, or 
                                    // - return 2 with probability 1/3. 
randomizedCollection.remove(1);     // return true since the collection contains 1. 
                                    // Removes 1 from the collection. Collection now contains [1,2]. 
randomizedCollection.getRandom();   // getRandom should return 1 or 2, both equally likely.

Constraints:

Solution

@Suppress("kotlin:S2245")
class RandomizedCollection() {
    val m2a = HashMap<Int, HashSet<Int>>()
    val a2m = ArrayList<Int>()
    /** Initialize your data structure here. */

    /** Inserts a value to the collection. Returns true if the collection did not already contain the specified element. */
    fun insert(x: Int): Boolean {
        a2m.add(x)
        val pos = a2m.size - 1
        return if (x in m2a) {
            m2a[x]!!.add(pos)
            false
        } else {
            m2a[x] = HashSet<Int>()
            m2a[x]!!.add(pos)
            true
        }
    }

    /** Removes a value from the collection. Returns true if the collection contained the specified element. */
    fun remove(x: Int): Boolean {
        if (x !in m2a) {
            return false
        }
        val pos = m2a[x]!!.iterator().next()
        if (m2a[x]!!.size == 1) {
            m2a.remove(x)
        } else {
            m2a[x]!!.remove(pos)
        }
        if (pos != a2m.size - 1) {
            m2a[a2m[a2m.size - 1]]!!.remove(a2m.size - 1)
            m2a[a2m[a2m.size - 1]]!!.add(pos)
            a2m[pos] = a2m[a2m.size - 1]
        }
        a2m.removeAt(a2m.size - 1)
        return true
    }

    /** Get a random element from the collection. */
    fun getRandom(): Int {
        val pos = Math.floor(Math.random() * a2m.size).toInt()
        return a2m[pos]
    }
}

/*
 * Your RandomizedCollection object will be instantiated and called as such:
 * var obj = RandomizedCollection()
 * var param_1 = obj.insert(`val`)
 * var param_2 = obj.remove(`val`)
 * var param_3 = obj.getRandom()
 */