LeetCode in Kotlin

1993. Operations on Tree

Medium

You are given a tree with n nodes numbered from 0 to n - 1 in the form of a parent array parent where parent[i] is the parent of the ith node. The root of the tree is node 0, so parent[0] = -1 since it has no parent. You want to design a data structure that allows users to lock, unlock, and upgrade nodes in the tree.

The data structure should support the following functions:

Implement the LockingTree class:

Example 1:

Input

[“LockingTree”, “lock”, “unlock”, “unlock”, “lock”, “upgrade”, “lock”]

[[[-1, 0, 0, 1, 1, 2, 2]], [2, 2], [2, 3], [2, 2], [4, 5], [0, 1], [0, 1]]

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

Explanation:

LockingTree lockingTree = new LockingTree([-1, 0, 0, 1, 1, 2, 2]);
lockingTree.lock(2, 2);    // return true because node 2 is unlocked.
                           // Node 2 will now be locked by user 2.
lockingTree.unlock(2, 3);  // return false because user 3 cannot unlock a node locked by user 2.
lockingTree.unlock(2, 2);  // return true because node 2 was previously locked by user 2.
                           // Node 2 will now be unlocked.
lockingTree.lock(4, 5);    // return true because node 4 is unlocked.
                           // Node 4 will now be locked by user 5.
lockingTree.upgrade(0, 1); // return true because node 0 is unlocked and has at least one locked descendant (node 4).
                           // Node 0 will now be locked by user 1 and node 4 will now be unlocked.
lockingTree.lock(0, 1);    // return false because node 0 is already locked. 

Constraints:

Solution

import java.util.LinkedList

class LockingTree(parent: IntArray) {
    private val a: Array<IntArray>
    private val map = HashMap<Int, MutableList<Int>>()

    init {
        val l = parent.size
        a = Array(l) { IntArray(2) }
        for (i in 0 until l) {
            a[i][0] = parent[i]
            a[i][1] = -1
            map.putIfAbsent(parent[i], ArrayList())
            val p = map[parent[i]]!!
            p.add(i)
            map[parent[i]] = p
        }
    }

    fun lock(num: Int, user: Int): Boolean {
        val userId = a[num][1]
        if (userId == -1) {
            a[num][1] = user
            return true
        }
        return false
    }

    fun unlock(num: Int, user: Int): Boolean {
        val y = a[num][1]
        if (y == user) {
            a[num][1] = -1
            return true
        }
        return false
    }

    fun upgrade(num: Int, user: Int): Boolean {
        var par = num
        while (par >= 0) {
            val lop = a[par][1]
            if (lop != -1) {
                return false
            }
            par = a[par][0]
        }
        var f = 0
        val que = LinkedList<Int>()
        val v = IntArray(a.size)
        que.add(num)
        v[num] = 1
        while (que.isNotEmpty()) {
            val t = que[0]
            que.removeAt(0)
            val p: List<Int> = map.getOrDefault(t, ArrayList())
            for (e in p) {
                if (a[e][1] != -1) {
                    f = 1
                    a[e][1] = -1
                }
                if (v[e] == 0) {
                    que.add(e)
                    v[e] = 1
                }
            }
        }
        if (f == 1) {
            a[num][1] = user
            return true
        }
        return false
    }
}
/*
 * Your LockingTree object will be instantiated and called as such:
 * var obj = LockingTree(parent)
 * var param_1 = obj.lock(num,user)
 * var param_2 = obj.unlock(num,user)
 * var param_3 = obj.upgrade(num,user)
 */