Problem Description

You need to implement a data structure for a tree that supports locking operations. The tree has n nodes numbered from 0 to n - 1, where node 0 is the root. The tree structure is given as a parent array where parent[i] indicates the parent of node i. Since node 0 is the root, parent[0] = -1.

The data structure must support three operations:

1. Lock Operation (lock(num, user))

• Locks node num for user user
• Returns true if locking is successful, false otherwise
• A node can only be locked if it's currently unlocked
• Once locked by a user, other users cannot lock the same node

2. Unlock Operation (unlock(num, user))

• Unlocks node num for user user
• Returns true if unlocking is successful, false otherwise
• A node can only be unlocked by the same user who locked it

3. Upgrade Operation (upgrade(num, user))

• Locks node num for user user AND unlocks all its descendants
• Returns true if upgrade is successful, false otherwise
• The upgrade can only happen when ALL three conditions are met:
  • The node num itself is unlocked
  • At least one descendant of num is locked (by any user)
  • No ancestor of num is locked

The solution uses:

• self.locked: An array tracking which user locked each node (-1 means unlocked)
• self.parent: The given parent array
• self.children: A list of children for each node (built from the parent array)

The upgrade operation uses DFS to traverse and unlock all descendants, while also checking if any descendants are locked. It traverses up the tree to verify no ancestors are locked before performing the upgrade.

Flowchart Walkthrough

First, let's pin down the algorithm using the Flowchart. Here's a step-by-step walkthrough:

Is it a graph?

• Yes: The problem involves a tree structure with nodes and edges (parent-child relationships). A tree is a special type of graph.

Is it a tree?

• Yes: The problem explicitly states we're working with a tree structure. We have n nodes with parent-child relationships defined by the parent array, and node 0 is the root.

DFS

• Conclusion: According to the flowchart, when we have a tree structure, the appropriate algorithm pattern is DFS (Depth-First Search).

Why DFS is the Right Choice

The flowchart correctly identifies DFS as the solution pattern because:

1. Tree Traversal: The upgrade operation requires traversing all descendants of a given node to unlock them. DFS is perfect for exploring all nodes in a subtree.

2. Recursive Nature: The tree structure naturally lends itself to recursive traversal, which is exactly what DFS provides. In the solution, the dfs function recursively visits each child and their descendants.

3. Complete Subtree Exploration: The upgrade operation needs to:

   • Check if any descendant is locked (requires visiting all descendants)
   • Unlock all locked descendants (requires accessing every node in the subtree)

   DFS ensures we visit every node in the subtree exactly once, making it efficient for this operation.

4. Parent-Child Relationships: The solution builds a children array to facilitate DFS traversal from any node downward through its descendants, which is essential for the upgrade operation.

The DFS pattern in the solution is implemented in the dfs helper function within the upgrade method, which recursively traverses all children and their descendants to perform the required unlocking operations.

Intuition

The key insight is recognizing that we need to efficiently track lock states and navigate both upward and downward in the tree structure.

For the basic lock and unlock operations, we simply need to track who has locked each node. A simple array where locked[i] stores the user ID (or -1 if unlocked) handles this elegantly.

The challenging part is the upgrade operation, which has three distinct requirements that need different traversal strategies:

1. Check no ancestors are locked: We need to traverse upward from the node to the root. Since we have the parent array, we can simply follow parent[x] repeatedly until we reach the root (parent[x] = -1). If we find any locked node along this path, the upgrade fails.

2. Check at least one descendant is locked AND unlock all descendants: These two requirements can be handled together with a single downward traversal. This is where DFS shines - we recursively visit all children and their descendants. During this traversal, we:

   • Track if we found any locked descendant using a flag variable
   • Unlock any locked nodes we encounter

The beauty of this approach is that we combine the checking and unlocking in a single DFS pass, making it efficient.

To enable downward traversal, we preprocess the parent array to build a children list for each node. This one-time preprocessing in the constructor allows us to quickly access all children of any node during DFS.

The solution elegantly uses:

• Upward traversal (following parent pointers) for ancestor checking
• Downward traversal (DFS through children) for descendant operations
• State tracking with the locked array for O(1) lock status checks

This combination of bidirectional tree navigation and simple state management provides an efficient solution to all three operations.

Learn more about Tree, Depth-First Search and Breadth-First Search patterns.

Solution Approach

Let's walk through the implementation of each component and operation in the LockingTree class:

Data Structure Initialization

The constructor sets up three key data structures:

1. self.locked: An array of size n initialized with -1, where locked[i] stores the user ID who locked node i, or -1 if unlocked.

2. self.parent: Stores the given parent array for upward traversal.

3. self.children: A list of lists built from the parent array. For each node i, children[i] contains all direct children of node i. This is constructed by iterating through the parent array starting from index 1 (skipping root) and appending each node to its parent's children list.

for son, fa in enumerate(parent[1:], 1):
    self.children[fa].append(son)

Lock Operation

The lock method is straightforward:

• Check if locked[num] == -1 (node is unlocked)
• If yes, set locked[num] = user and return True
• Otherwise, return False

Time complexity: O(1)

Unlock Operation

The unlock method verifies ownership:

• Check if locked[num] == user (same user who locked it)
• If yes, set locked[num] = -1 and return True
• Otherwise, return False

Time complexity: O(1)

Upgrade Operation

The upgrade method is the most complex, implementing three checks:

Step 1: Check no ancestors are locked

x = num
while x != -1:
    if self.locked[x] != -1:
        return False
    x = self.parent[x]

Starting from num, traverse upward to the root. If any ancestor is locked, immediately return False.

Step 2: Find and unlock all locked descendants

def dfs(x: int):
    nonlocal find
    for y in self.children[x]:
        if self.locked[y] != -1:
            self.locked[y] = -1
            find = True
        dfs(y)

The DFS helper function:

• Recursively visits all children and their descendants
• When a locked node is found, unlocks it (locked[y] = -1) and sets find = True
• The nonlocal find allows the nested function to modify the outer scope's find variable

Step 3: Finalize the upgrade

find = False
dfs(num)
if not find:
    return False
self.locked[num] = user
return True

• Initialize find to False before DFS
• Run DFS to unlock descendants and detect if any were locked
• If no locked descendants were found (find is still False), return False
• Otherwise, lock the node for the user and return True

Time Complexity Analysis

• Lock/Unlock: O(1) - Simple array access and update
• Upgrade: O(h + d) where h is the height from node to root (for ancestor check) and d is the number of descendants (for DFS traversal)

The solution efficiently handles all operations by maintaining proper data structures for both upward and downward tree traversal, making it practical even for large trees.

Example Walkthrough

Let's walk through a concrete example with a small tree to illustrate how the locking operations work.

Tree Structure:

parent = [-1, 0, 0, 1, 1, 3]

Tree visualization:
        0
       / \
      1   2
     / \
    3   4
   /
  5

Initial Setup:

• locked = [-1, -1, -1, -1, -1, -1] (all nodes unlocked)
• children = [[1,2], [3,4], [], [5], [], []]

Operation Sequence:

1. lock(1, 10) → Returns true

   • Check: locked[1] == -1? Yes, node 1 is unlocked
   • Action: Set locked[1] = 10
   • State: locked = [-1, 10, -1, -1, -1, -1]

2. lock(3, 20) → Returns true

   • Check: locked[3] == -1? Yes, node 3 is unlocked
   • Action: Set locked[3] = 20
   • State: locked = [-1, 10, -1, 20, -1, -1]

3. unlock(3, 10) → Returns false

   • Check: locked[3] == 10? No, it's locked by user 20
   • Action: No change
   • State: locked = [-1, 10, -1, 20, -1, -1]

4. upgrade(0, 30) → Returns false

   • Check ancestors of node 0: Node 0 has no ancestors (it's the root)
   • Check descendants: Found locked nodes 1 and 3
   • But node 0's direct child (node 1) is locked, violating the "node itself must be unlocked" condition
   • Action: No change

5. unlock(1, 10) → Returns true

   • Check: locked[1] == 10? Yes
   • Action: Set locked[1] = -1
   • State: locked = [-1, -1, -1, 20, -1, -1]

6. upgrade(1, 40) → Returns true

   • Check ancestors: Traverse 1→0, check if any are locked
     • locked[1] == -1? Yes (unlocked)
     • locked[0] == -1? Yes (unlocked)
   • Check descendants via DFS from node 1:
     • Visit child 3: locked[3] == 20 (locked!) → Unlock it, set find = true
     • Visit child 4: locked[4] == -1 (already unlocked)
     • Visit grandchild 5: locked[5] == -1 (already unlocked)
   • Since find == true (found at least one locked descendant), proceed
   • Action: Lock node 1 for user 40
   • Final state: locked = [-1, 40, -1, -1, -1, -1]

7. upgrade(3, 50) → Returns false

   • Check ancestors: Traverse 3→1→0
     • locked[3] == -1? Yes (unlocked)
     • locked[1] == 40? No! Node 1 is locked
   • Since an ancestor is locked, immediately return false
   • Action: No change

This example demonstrates all three operations and shows how the upgrade operation carefully checks both ancestors (upward traversal) and descendants (DFS traversal) while managing the lock states.

Solution Implementation

Time and Space Complexity

Time Complexity:

• __init__(parent): O(n) where n is the length of the parent array. We iterate through the parent array once to build the children adjacency list.

• lock(num, user): O(1). We only check and update a single element in the locked array.

• unlock(num, user): O(1). We only check and update a single element in the locked array.

• upgrade(num, user): O(n) in the worst case.

  • Checking ancestors: O(h) where h is the height of the tree. In the worst case (skewed tree), h = n - 1.
  • DFS to unlock descendants: O(m) where m is the number of nodes in the subtree rooted at num. In the worst case, this could be O(n) if most nodes are descendants.
  • Overall: O(h + m) = O(n) in the worst case.

Space Complexity:

• __init__(parent): O(n) for storing:

  • self.locked: O(n) array
  • self.parent: O(n) array
  • self.children: O(n) for the adjacency list (total edges = n - 1)

• lock(num, user): O(1). No additional space used.

• unlock(num, user): O(1). No additional space used.

• upgrade(num, user): O(h) where h is the height of the tree, due to the recursive call stack in the DFS function. In the worst case (skewed tree), this could be O(n).

Overall Space Complexity of the class: O(n) for the data structures maintained throughout the object's lifetime.

Learn more about how to find time and space complexity quickly.

Common Pitfalls

1. Incorrect Ancestor Check in Upgrade Operation

A critical pitfall is checking ancestors incorrectly during the upgrade operation. Many implementations mistakenly start checking from the parent of the target node instead of the node itself.

Incorrect Implementation:

def upgrade(self, num: int, user: int) -> bool:
    # WRONG: Starting from parent, missing the check for num itself
    current_node = self.parent[num]  # Skips checking if num is locked!
    while current_node != -1:
        if self.locked[current_node] != -1:
            return False
        current_node = self.parent[current_node]

Why This Fails: According to the problem requirements, the upgrade operation requires that "the node num itself is unlocked." If we start checking from the parent, we miss verifying whether the target node is already locked, which would violate the upgrade conditions.

Correct Implementation:

def upgrade(self, num: int, user: int) -> bool:
    # CORRECT: Start from the node itself
    current_node = num  # Include num in the ancestor check
    while current_node != -1:
        if self.locked[current_node] != -1:
            return False
        current_node = self.parent[current_node]

2. Forgetting to Use nonlocal in Nested DFS Function

When using a nested function to track whether any descendants are locked, forgetting the nonlocal keyword causes the variable update to be local to the nested function only.

Incorrect Implementation:

def upgrade(self, num: int, user: int) -> bool:
    has_locked_descendant = False
  
    def unlock_descendants(node: int) -> None:
        # WRONG: Missing nonlocal, creates a new local variable
        for child in self.children[node]:
            if self.locked[child] != -1:
                self.locked[child] = -1
                has_locked_descendant = True  # This creates a LOCAL variable!
            unlock_descendants(child)
  
    unlock_descendants(num)
    # has_locked_descendant is still False here, even if descendants were locked
    if not has_locked_descendant:
        return False

Correct Implementation:

def upgrade(self, num: int, user: int) -> bool:
    has_locked_descendant = False
  
    def unlock_descendants(node: int) -> None:
        nonlocal has_locked_descendant  # CORRECT: Modifies outer scope variable
        for child in self.children[node]:
            if self.locked[child] != -1:
                self.locked[child] = -1
                has_locked_descendant = True
            unlock_descendants(child)

3. Inefficient Children List Construction

Building the children adjacency list incorrectly can lead to index errors or missing relationships.

Incorrect Implementation:

def __init__(self, parent: List[int]):
    self.children = [[] for _ in range(len(parent))]
    # WRONG: Starting from 0 includes root, which has parent -1
    for child_node, parent_node in enumerate(parent):
        if parent_node != -1:  # Extra check needed
            self.children[parent_node].append(child_node)

Better Implementation:

def __init__(self, parent: List[int]):
    self.children = [[] for _ in range(len(parent))]
    # CORRECT: Start from index 1, skip root automatically
    for child_node, parent_node in enumerate(parent[1:], 1):
        self.children[parent_node].append(child_node)

The second approach is cleaner as it naturally skips the root node (which has no parent) and avoids the need for conditional checks.

4. Not Handling Edge Cases Properly

Failing to consider edge cases like single-node trees or operations on the root node.

Example Edge Case:

• Tree with only root node (n=1): The upgrade operation should always fail since root has no descendants
• Attempting to upgrade the root when it's the only locked node in the tree

Solution: The implementation should naturally handle these cases through the DFS traversal, which will find no descendants for a leaf node, causing the upgrade to correctly return False.