Problem

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:

  • Lock: Locks the given node for the given user and prevents other users from locking the same node. You may only lock a node using this function if the node is unlocked.
  • Unlock: Unlocks the given node for the given user. You may only unlock a node using this function if it is currently locked by the same user.
  • Upgrade**: Locks** the given node for the given user and unlocks all of its descendants regardless of who locked it. You may only upgrade a node if all 3 conditions are true:
    • The node is unlocked,
    • It has at least one locked descendant (by any user), and
    • It does not have any locked ancestors.

Implement the LockingTree class:

  • LockingTree(int[] parent) initializes the data structure with the parent array.
  • lock(int num, int user) returns true if it is possible for the user with id user to lock the node num, or false otherwise. If it is possible, the node num will become locked by the user with id user.
  • unlock(int num, int user) returns true if it is possible for the user with id user to unlock the node num, or false otherwise. If it is possible, the node num will become unlocked.
  • upgrade(int num, int user) returns true if it is possible for the user with id user to upgrade the node num, or false otherwise. If it is possible, the node num will be upgraded.

Examples

Example 1:

graph TD;
	A(0) --- B(1) & C(2)
	B --- D(3) & E(4)
	C --- F(5) & G(6)
  
 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18

**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.

Solution

Method 1

Code

 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
71
72
73
74
75
76
77
78
79
80
81
82
83

class LockingTree {
	private int[] parent;
	private Map<Integer, Integer> locked; // {node: user}
	private Map<Integer, List<Integer>> children;
	
	public LockingTree(int[] parent) {
		this.parent = parent;
		this.locked = new HashMap<>();
		this.children = new HashMap<>();
		for (int i = 0; i < parent.length; i++) {
			children.put(i, new ArrayList<>());
		}
		for (int i = 0; i < parent.length; i++) {
			if (parent[i] != -1) {
				children.get(parent[i]).add(i);
			}
		}
	}
	
	public boolean lock(int num, int user) {
		if (!locked.containsKey(num)) {
			locked.put(num, user);
			return true;
		}
		return false;
	}
	
	public boolean unlock(int num, int user) {
		if (locked.containsKey(num) && locked.get(num) == user) {
			locked.remove(num);
			return true;
		}
		return false;
	}
	
	public boolean upgrade(int num, int user) {
		// Check if num is unlocked
		if (locked.containsKey(num)) {
			return false;
		}
		
		// Check if ancestors are not locked
		int cur = num;
		while (cur != -1) {
			if (locked.containsKey(cur)) {
				return false;
			}
			cur = parent[cur];
		}
		
		// Check if there is at least one locked descendant
		boolean hasLockedDescendant = hasLockedDescendant(num);
		if (!hasLockedDescendant) {
			return false;
		}
		
		// Unlock all descendants
		unlockDescendants(num);
		locked.put(num, user);
		return true;
	}
	
	private boolean hasLockedDescendant(int num) {
		if (locked.containsKey(num)) {
			return true;
		}
		for (int child : children.get(num)) {
			if (hasLockedDescendant(child)) {
				return true;
			}
		}
		return false;
	}
	
	private void unlockDescendants(int num) {
		if (locked.containsKey(num)) {
			locked.remove(num);
		}
		for (int child : children.get(num)) {
			unlockDescendants(child);
		}
	}
}
 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

class LockingTree:
    def __init__(self, parent: List[int]):
        self.parent = parent
        self.locked: Dict[int, int] = {}  # {node: user}
        self.children = {i: [] for i in range(len(parent))}
        for i in range(len(parent)):
            if parent[i] != -1:
                self.children[parent[i]].append(i)
    
    def lock(self, num: int, user: int) -> bool:
        if num not in self.locked:
            self.locked[num] = user
            return True
        return False
    
    def unlock(self, num: int, user: int) -> bool:
        if num in self.locked and self.locked[num] == user:
            del self.locked[num]
            return True
        return False
    
    def upgrade(self, num: int, user: int) -> bool:
        # Check if num is unlocked
        if num in self.locked:
            return False
        
        # Check if ancestors are not locked
        cur = num
        while cur != -1:
            if cur in self.locked:
                return False
            cur = self.parent[cur]
        
        # Check if there is at least one locked descendant
        def has_locked_descendant(node: int) -> bool:
            if node in self.locked:
                return True
            for child in self.children[node]:
                if has_locked_descendant(child):
                    return True
            return False
        
        if not has_locked_descendant(num):
            return False
        
        # Unlock all descendants
        def unlock_descendants(node: int):
            if node in self.locked:
                del self.locked[node]
            for child in self.children[node]:
                unlock_descendants(child)
        
        unlock_descendants(num)
        self.locked[num] = user
        return True

Complexity

  • ⏰ Time complexity:
    • lock/unlock: O(1) — Both operations are simple hash map lookups and updates.
    • upgrade: O(h + d)
      • Checking ancestors (to root): O(h), where h is the height of the tree.
      • Checking for locked descendants and unlocking them: O(d), where d is the number of descendants (in the worst case, O(n) if the tree is a chain or the node is the root).
      • So, worst-case for upgrade is O(n).
  • 🧺 Space complexity:
    • The parent array is O(n).
    • The locked map can have up to O(n) entries.
    • The children map/list is O(n) for storing all children relationships.