Operations on Tree: Lock, Unlock, and Upgrade

You are given a tree with n nodes numbered 0 to n - 1, rooted at node 0, and represented as a parent array. Design a data structure LockingTree that supports three operations: lock(num, user) locks node num by user if unlocked, unlock(num, user) unlocks node num only if locked by that same user, and upgrade(num, user) locks node num and unlocks all of its locked descendants, but only if node num is currently unlocked, it has at least one locked descendant, and none of its ancestors are locked.

This is LeetCode 1993: Operations on Tree, and it tests your ability to combine tree traversal techniques with clean class design. The parent array gives you upward traversal for free (just follow the chain), while DFS handles downward traversal through descendants.

Why this problem matters

This problem sits at the intersection of two skills that come up constantly in interviews: tree traversal and data structure design. The lock and unlock operations are simple dictionary lookups, but the upgrade operation forces you to think carefully about how to traverse both up (ancestors) and down (descendants) efficiently.

It also tests whether you can translate a multi-condition specification into clean code. The upgrade operation has three conditions that must all be satisfied. Missing any one of them means a wrong answer. Problems like this reward careful reading and methodical implementation, not clever tricks.

The approach

The key insight is that the parent array makes ancestor traversal trivial. To check if any ancestor is locked, just walk up from parent[num] to parent[parent[num]] and so on until you reach -1 (the root's parent). Each step is O(1), and the total walk is at most O(n) in a skewed tree but typically O(log n) in a balanced one.

For descendants, build an adjacency list (children mapping) from the parent array during initialization. Then use DFS from node num to visit all descendants. During the DFS, collect any locked nodes and unlock them.

Lock and unlock are just dictionary operations. Store a locks dictionary mapping node_number to user_id. Locking checks if the key exists. Unlocking checks if the key exists AND the value matches the user.

Clean Python solution

class LockingTree:
    def __init__(self, parent):
        self.parent = parent
        self.locks = {}
        self.children = [[] for _ in range(len(parent))]
        for i in range(1, len(parent)):
            self.children[parent[i]].append(i)

    def lock(self, num, user):
        if num in self.locks:
            return False
        self.locks[num] = user
        return True

    def unlock(self, num, user):
        if self.locks.get(num) != user:
            return False
        del self.locks[num]
        return True

    def upgrade(self, num, user):
        if num in self.locks:
            return False

        cur = self.parent[num]
        while cur != -1:
            if cur in self.locks:
                return False
            cur = self.parent[cur]

        locked_descendants = []
        stack = [num]
        while stack:
            node = stack.pop()
            for child in self.children[node]:
                if child in self.locks:
                    locked_descendants.append(child)
                stack.append(child)

        if not locked_descendants:
            return False

        for node in locked_descendants:
            del self.locks[node]
        self.locks[num] = user
        return True

Step-by-step walkthrough

The critical step is Step 5. The upgrade operation has three checks, and all three must pass before anything changes. First, verify the target node is unlocked. Second, walk up the parent array to confirm no ancestor is locked. Third, DFS down through descendants to find at least one locked node. Only after all three checks pass do you unlock the descendants and lock the target.

Notice how the DFS collects all locked descendants into a list before modifying anything. This avoids issues with modifying the lock dictionary while you are still searching through the tree. Clean separation between the "check" phase and the "modify" phase makes the code easier to reason about.

Complexity analysis

lock and unlock are O(1) dictionary operations. The space is O(n) for the lock dictionary in the worst case where every node is locked.

upgrade is O(n) in the worst case. Walking up to the root costs O(depth), which is O(n) for a skewed tree. The DFS through all descendants can visit every node. Combined, the upgrade operation is O(n) time. The DFS stack and the locked descendants list use O(n) space.

The building blocks

This problem combines three reusable pieces that CodeBricks drills independently:

1. Parent array traversal

The parent array parent[i] gives you the parent of node i. To walk up the tree:

cur = parent[num]
while cur != -1:
    # process ancestor cur
    cur = parent[cur]

This is a linear chain walk, identical to traversing a linked list. You follow the pointer (parent reference) one step at a time until you hit the sentinel value -1. The same pattern appears in Union-Find (following parent pointers to find the root) and in any tree problem that gives you the parent array instead of an adjacency list.

2. DFS over children (descendants)

To visit all descendants, build a children list from the parent array and DFS:

children = [[] for _ in range(n)]
for i in range(1, n):
    children[parent[i]].append(i)

stack = [start]
while stack:
    node = stack.pop()
    for child in children[node]:
        stack.append(child)

This is standard iterative DFS over a tree. The same pattern shows up in Binary Tree Level Order Traversal (with a queue instead of a stack) and in any tree problem where you need to process all nodes in a subtree.

3. Dictionary-based state tracking

Using a dictionary to track which nodes are locked and by whom:

locks = {}
locks[num] = user      # lock
del locks[num]          # unlock
num in locks            # is locked?
locks.get(num) == user  # locked by this user?

This is the same "dictionary as state" pattern you see in LRU Cache and other design problems. The dictionary gives you O(1) lookups, insertions, and deletions. It is also space-efficient because you only store entries for locked nodes rather than maintaining a flag on every node.

Edge cases

• Single node tree: parent = [-1]. Lock and unlock work normally. Upgrade always fails because a single node has no descendants.
• Locking an already locked node: lock(num, user) returns False if num is already in the locks dictionary, regardless of the user.
• Wrong user unlock: unlock(num, user) returns False if the node is locked by a different user. The stored user must match exactly.
• Upgrade with no locked descendants: Even if the target is unlocked and has no locked ancestors, the upgrade fails if there are zero locked descendants. The "at least one" condition is required.
• Upgrade with a locked ancestor: If any node on the path from num to the root is locked, upgrade fails. You must walk the entire path up to confirm.
• Empty locks dictionary: When no nodes are locked, lock always succeeds, unlock always fails, and upgrade always fails (no locked descendants).

From understanding to recall

You have read through the solution and traced the steps. You see how the parent array handles upward traversal, how DFS handles downward traversal, and how the dictionary tracks lock state. The logic makes sense. But can you write the entire LockingTree class from scratch in five minutes during an interview?

The details trip people up. Building the children list in the constructor. Using self.locks.get(num) != user instead of checking existence and value separately. Collecting locked descendants before modifying the dictionary. These are not conceptually hard, but they are exactly the kind of things you fumble under time pressure if you have not drilled them.

Spaced repetition makes these details automatic. You practice writing the parent traversal, the DFS over children, and the three-condition upgrade check from memory at increasing intervals. After a few rounds, the pattern flows out without hesitation. You do not need to re-derive it. You just write it.

Related posts

• Clone Graph - Another tree/graph structure problem requiring careful traversal
• Binary Tree Level Order Traversal - Tree traversal patterns that show up in lock checking
• Design Add and Search Words - Another design problem combining trees with specific operations

CodeBricks breaks the Operations on Tree problem into its parent-traversal, DFS-over-children, and dictionary-state-tracking building blocks, then drills them independently with spaced repetition. You type each piece from scratch until the pattern is automatic. When a tree design problem shows up in your interview, you do not think about it. You just write it.