LeetCode in Kotlin

1172. Dinner Plate Stacks

Hard

You have an infinite number of stacks arranged in a row and numbered (left to right) from 0, each of the stacks has the same maximum capacity.

Implement the DinnerPlates class:

Example 1:

Input

[“DinnerPlates”, “push”, “push”, “push”, “push”, “push”, “popAtStack”, “push”, “push”, “popAtStack”, “popAtStack”, “pop”, “pop”, “pop”, “pop”, “pop”]

[[2], [1], [2], [3], [4], [5], [0], [20], [21], [0], [2], [], [], [], [], []]

Output: [null, null, null, null, null, null, 2, null, null, 20, 21, 5, 4, 3, 1, -1]

Explanation:

DinnerPlates D = DinnerPlates(2); // Initialize with capacity = 2
D.push(1); 
D.push(2); 
D.push(3); 
D.push(4); 
D.push(5);          // The stacks are now: 2 4 
                                          1 3 5 
                                          ﹈ ﹈ ﹈ 
D.popAtStack(0);    // Returns 2. The stacks are now:   4 
                                                   1 3 5 
                                                  ﹈ ﹈ ﹈ 
D.push(20);         // The stacks are now: 20 4 
                                            1 3 5 
                                           ﹈ ﹈ ﹈ 
D.push(21);         // The stacks are now: 20 4 21 
                                            1 3 5 
                                            ﹈ ﹈ ﹈ 
D.popAtStack(0);    // Returns 20. The stacks are now: 4 21 
                                                       1 3 5 
                                                      ﹈ ﹈ ﹈ 
D.popAtStack(2);    // Returns 21. The stacks are now: 4 
                                                     1 3 5 
                                                    ﹈ ﹈ ﹈ 
D.pop()             // Returns 5. The stacks are now: 4 
                                                     1 3 
                                                     ﹈ ﹈ 
D.pop()             // Returns 4. The stacks are now: 1 3 
                                                     ﹈ ﹈ 
D.pop()             // Returns 3. The stacks are now: 1 
                                                      ﹈ 
D.pop()             // Returns 1. There are no stacks. 
D.pop()             // Returns -1. There are still no stacks.

Constraints:

Solution

import java.util.TreeSet

class DinnerPlates(private val stackCap: Int) {
    private val stacks: MutableList<ArrayDeque<Int>>
    private val leftIndex: TreeSet<Int>

    init {
        stacks = ArrayList()
        leftIndex = TreeSet()
    }

    fun push(`val`: Int) {
        if (leftIndex.isNotEmpty()) {
            val i = leftIndex.first()
            stacks[i].addLast(`val`)
            if (stacks[i].size == stackCap) {
                leftIndex.remove(i)
            }
            return
        }
        if (stacks.isEmpty() || stacks[stacks.size - 1].size == stackCap) {
            val newStack = ArrayDeque<Int>()
            stacks.add(newStack)
        }
        stacks[stacks.size - 1].addLast(`val`)
    }

    fun pop(): Int {
        if (stacks.isEmpty()) {
            return -1
        }
        while (stacks[stacks.size - 1].isEmpty()) {
            leftIndex.remove(stacks.size - 1)
            stacks.removeAt(stacks.size - 1)
        }
        val `val` = stacks[stacks.size - 1].removeLast()
        if (stacks[stacks.size - 1].isEmpty()) {
            leftIndex.remove(stacks.size - 1)
            stacks.removeAt(stacks.size - 1)
        }
        return `val`
    }

    fun popAtStack(index: Int): Int {
        if (stacks.size - 1 >= index) {
            var `val` = -1
            if (stacks[index].isNotEmpty()) {
                `val` = stacks[index].removeLast()
            }
            if (stacks[index].isEmpty() && index == stacks.size - 1) {
                leftIndex.remove(stacks.size - 1)
                stacks.removeAt(index)
                return `val`
            }
            leftIndex.add(index)
            return `val`
        }
        return -1
    }
}