LeetCode Solutions Blog

Operations on Tree

Number: 2104

Difficulty: Medium

Paid? No

Problem Description

We are given a tree with n nodes (numbered from 0 to n-1) represented via a parent array, where parent[0] = -1 for the root node. We need to implement a LockingTree class that supports three operations:

lock(num, user): Lock the node num for the given user if it is currently unlocked.
unlock(num, user): Unlock node num if it is currently locked by the same user.
upgrade(num, user): If node num is unlocked, has at least one locked descendant, and none of its ancestors are locked, then lock node num for the user and unlock all of its descendants.

Key Insights

Use an array (or similar data structure) to keep track of the lock status of each node (e.g., using -1 for unlocked and a valid user id when locked).
Pre-compute a children mapping from the parent array, allowing efficient DFS/BFS traversal when checking or unlocking descendants.
When upgrading, check upward (using the parent pointers) to ensure no locked ancestors.
Use DFS from the node to both check for any locked descendants and to unlock them if needed.

Space and Time Complexity

Time Complexity:

lock and unlock operations: O(1)
upgrade operation: O(n) in worst-case due to traversal of the descendant subtree.

Space Complexity: O(n) for storing the children list and the lock state array.

Solution

The solution involves:

Building a children list from the parent array to efficiently traverse descendants.
Maintaining an array for the locked status of each node.
For the lock operation, directly update the state if the node is unlocked.
For the unlock operation, verify the user before unlocking.
For the upgrade operation, first check that the node is unlocked, then traverse up the tree to ensure there are no locked ancestors. Next, perform a DFS to check if there is any locked descendant while unlocking them. If at least one descendant was locked, the current node is locked with the provided user.

Code Solutions

class LockingTree:
    def __init__(self, parent):
        # Store the parent pointers
        self.parent = parent
        n = len(parent)
        # -1 indicates the node is unlocked; otherwise, store the user id who locked it
        self.locked = [-1] * n
        # Build children list for DFS traversal
        self.children = [[] for _ in range(n)]
        for i in range(1, n):
            self.children[parent[i]].append(i)

    def lock(self, num, user):
        # If already locked, cannot lock it again
        if self.locked[num] != -1:
            return False
        self.locked[num] = user
        return True

    def unlock(self, num, user):
        # Unlock only if the node is locked by the same user
        if self.locked[num] != user:
            return False
        self.locked[num] = -1
        return True

    def upgrade(self, num, user):
        # Condition 1: the node must be unlocked.
        if self.locked[num] != -1:
            return False
        # Condition 3: ensure no locked ancestors.
        curr = self.parent[num]
        while curr != -1:
            if self.locked[curr] != -1:
                return False
            curr = self.parent[curr]
        
        # Condition 2: there must be at least one locked descendant.
        # DFS to find any locked descendant and unlock them.
        locked_found = False
        def dfs(node):
            nonlocal locked_found
            if self.locked[node] != -1:
                self.locked[node] = -1  # Unlock this descendant
                locked_found = True
            for child in self.children[node]:
                dfs(child)
        
        dfs(num)
        
        if not locked_found:
            return False
        
        # Lock current node since conditions are met.
        self.locked[num] = user
        return True

# Example usage:
# tree = LockingTree([-1, 0, 0, 1, 1, 2, 2])
# print(tree.lock(2, 2))      # True
# print(tree.unlock(2, 3))    # False
# print(tree.unlock(2, 2))    # True
# print(tree.lock(4, 5))      # True
# print(tree.upgrade(0, 1))   # True (locks node 0 and unlocks its locked descendant(s))
# print(tree.lock(0, 1))      # False