# Operations on Tree
**Difficulty:** MEDIUM
[External](https://leetcode.com/problems/operations-on-tree)
Canonical: https://scaleengineer.com/dsa/problems/operations-on-tree
**Patterns:** [Design](https://scaleengineer.com/dsa/patterns/design)
**Algorithms:** [Depth-First Search](https://scaleengineer.com/algorithms/depth-first-search), [Breadth-First Search](https://scaleengineer.com/algorithms/breadth-first-search)
**Data structures:** Array, Hash Table, Tree
**Companies:** [Juspay](https://scaleengineer.com/companies/juspay)
---
## 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**.

**Example 1:**

![](https://assets.glich.co/dsa/operations-on-tree/image0.png) 

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

* `n == parent.length`
* `2 <= n <= 2000`
* `0 <= parent[i] <= n - 1` for `i != 0`
* `parent[0] == -1`
* `0 <= num <= n - 1`
* `1 <= user <= 104`
* `parent` represents a valid tree.
* At most `2000` calls **in total** will be made to `lock`, `unlock`, and `upgrade`.

# Approaches
## Tree Traversal with Adjacency List
This approach focuses on creating an efficient data structure to handle the three required operations. We pre-process the parent array to build an adjacency list representation of the tree, which allows for efficient downward traversal (to find descendants). The lock status of each node is tracked in a separate array.

- The `lock` and `unlock` operations are straightforward, involving a simple check and update on the lock status array, making them very fast (O(1)).
- The `upgrade` operation is the most complex. It validates three conditions:
  1. The node itself is unlocked.
  2. None of its ancestors are locked (checked by traversing up the `parent` array).
  3. At least one of its descendants is locked (checked by traversing down using the pre-built adjacency list).

If all conditions pass, the node is locked, and all its locked descendants are found and unlocked.
**Time:** - **Constructor:** O(N) to build the adjacency list.
- **`lock`:** O(1).
- **`unlock`:** O(1).
- **`upgrade`:** O(H + D), where H is the height of the tree and D is the number of descendants of the node. In the worst case (a skewed tree or a node near the root), this is O(N). · **Space:** O(N), where N is the number of nodes. This is used to store the `parent` array, the `lockedBy` status array, and the `children` adjacency list. The `upgrade` function may also use up to O(N) space for the traversal queue and the list of locked descendants.
**Pros:** The `lock` and `unlock` operations are highly efficient, with a constant time complexity of O(1).; Pre-computation of the children list in the constructor optimizes the repeated need for descendant traversal in the `upgrade` function.; The logic is straightforward and directly maps to the problem's requirements.
**Cons:** The `upgrade` operation has a time complexity of O(N) in the worst case, which could be slow for very large trees, although it is acceptable for the given constraints.
### Explanation
### Data Structures
- `int[] parent`: Stores the parent of each node.
- `int[] lockedBy`: An array of size `n` where `lockedBy[i]` stores the user ID that has locked node `i`, or `0` if it's unlocked.
- `List<Integer>[] children`: An adjacency list where `children[i]` contains a list of all direct children of node `i`. This is built once in the constructor for efficient descendant traversal.

### Constructor `LockingTree(int[] parent)`
The constructor initializes the data structures. It populates the `children` adjacency list by iterating through the `parent` array. This O(N) pre-processing step avoids re-computing child relationships on every `upgrade` call.
```java
class LockingTree {
    private int[] parent;
    private int[] lockedBy; // 0 if unlocked, user ID otherwise
    private List<Integer>[] children;

    public LockingTree(int[] parent) {
        int n = parent.length;
        this.parent = parent;
        this.lockedBy = new int[n];
        this.children = new ArrayList[n];
        for (int i = 0; i < n; i++) {
            children[i] = new ArrayList<>();
        }
        for (int i = 1; i < n; i++) {
            children[parent[i]].add(i);
        }
    }
```

### `lock(int num, int user)` and `unlock(int num, int user)`
These methods are simple O(1) operations that check and modify the `lockedBy` array.
```java
    public boolean lock(int num, int user) {
        if (lockedBy[num] == 0) {
            lockedBy[num] = user;
            return true;
        }
        return false;
    }

    public boolean unlock(int num, int user) {
        if (lockedBy[num] == user) {
            lockedBy[num] = 0;
            return true;
        }
        return false;
    }
```

### `upgrade(int num, int user)`
This method implements the three-condition check. It first checks the node's own lock status. Then, it walks up the tree to check for locked ancestors. Finally, it uses BFS to traverse the subtree and find locked descendants. If all conditions are satisfied, it performs the lock and unlock operations.
```java
    public boolean upgrade(int num, int user) {
        // Condition 1: The node is unlocked.
        if (lockedBy[num] != 0) {
            return false;
        }

        // Condition 3: It does not have any locked ancestors.
        int curr = num;
        while (parent[curr] != -1) {
            curr = parent[curr];
            if (lockedBy[curr] != 0) {
                return false;
            }
        }

        // Condition 2: It has at least one locked descendant.
        List<Integer> lockedDescendants = new ArrayList<>();
        Queue<Integer> queue = new LinkedList<>();
        queue.add(num);

        int count = 0;
        while (!queue.isEmpty()) {
            int node = queue.poll();
            if(count > 0 && lockedBy[node] != 0) { // count > 0 to skip the node `num` itself
                lockedDescendants.add(node);
            }
            count++;
            for (int child : children[node]) {
                queue.add(child);
            }
        }

        if (lockedDescendants.isEmpty()) {
            return false;
        }

        // All conditions met, perform the upgrade.
        lockedBy[num] = user;
        for (int descendant : lockedDescendants) {
            lockedBy[descendant] = 0;
        }

        return true;
    }
}
```
### Algorithm
1.  **Initialization (`LockingTree` constructor):**
    *   Store the input `parent` array.
    *   Initialize a `lockedBy` array of size `n` with `0`s, where `0` signifies an unlocked node.
    *   Create an adjacency list `children` to store the tree structure in a way that's easy to traverse downwards. Iterate through the `parent` array from node `1` to `n-1` and for each node `i`, add it to the children list of `parent[i]`.
2.  **`lock(num, user)`:**
    *   Check if `lockedBy[num]` is `0` (unlocked).
    *   If it is, update `lockedBy[num] = user` and return `true`.
    *   Otherwise, return `false`.
3.  **`unlock(num, user)`:**
    *   Check if `lockedBy[num]` is equal to `user`.
    *   If it is, update `lockedBy[num] = 0` and return `true`.
    *   Otherwise, return `false`.
4.  **`upgrade(num, user)`:**
    *   **Condition 1 (Node is Unlocked):** Check if `lockedBy[num]` is `0`. If not, return `false` immediately.
    *   **Condition 3 (No Locked Ancestors):** Traverse upwards from `num` using the `parent` array. For each ancestor, check its lock status in `lockedBy`. If any ancestor is locked, return `false`.
    *   **Condition 2 (Has Locked Descendants):** Perform a traversal (like Breadth-First Search or Depth-First Search) starting from `num`'s children to find all its descendants. 
        *   Use a queue for BFS and a list to store any locked descendants found.
        *   If the traversal completes and no locked descendants were found, return `false`.
    *   **Execution:** If all three conditions are met, proceed with the upgrade:
        *   Lock the current node: `lockedBy[num] = user`.
        *   Iterate through the list of locked descendants collected earlier and unlock each one by setting their `lockedBy` entry to `0`.
        *   Return `true`.

# Solutions
### Java

```java
class LockingTree { private int [] locked ; private int [] parent ; private List < Integer >[] children ; public LockingTree ( int [] parent ) { int n = parent . length ; locked = new int [ n ]; this . parent = parent ; children = new List [ n ]; Arrays . fill ( locked , - 1 ); Arrays . setAll ( children , i -> new ArrayList <>()); for ( int i = 1 ; i < n ; i ++) { children [ parent [ i ]]. add ( i ); } } public boolean lock ( int num , int user ) { if ( locked [ num ] == - 1 ) { locked [ num ] = user ; return true ; } return false ; } public boolean unlock ( int num , int user ) { if ( locked [ num ] == user ) { locked [ num ] = - 1 ; return true ; } return false ; } public boolean upgrade ( int num , int user ) { int x = num ; while ( x != - 1 ) { if ( locked [ x ] != - 1 ) { return false ; } x = parent [ x ]; } boolean [] find = new boolean [ 1 ]; dfs ( num , find ); if (! find [ 0 ]) { return false ; } locked [ num ] = user ; return true ; } private void dfs ( int x , boolean [] find ) { for ( int y : children [ x ]) { if ( locked [ y ] != - 1 ) { locked [ y ] = - 1 ; find [ 0 ] = true ; } dfs ( y , find ); } } } /** * Your LockingTree object will be instantiated and called as such: * LockingTree obj = new LockingTree(parent); * boolean param_1 = obj.lock(num,user); * boolean param_2 = obj.unlock(num,user); * boolean param_3 = obj.upgrade(num,user); */
```

### CPP

```cpp
class LockingTree { public: LockingTree ( vector < int >& parent ) { int n = parent . size (); locked = vector < int > ( n , - 1 ); this -> parent = parent ; children . resize ( n ); for ( int i = 1 ; i < n ; ++ i ) { children [ parent [ i ]]. push_back ( i ); } } bool lock ( int num , int user ) { if ( locked [ num ] == - 1 ) { locked [ num ] = user ; return true ; } return false ; } bool unlock ( int num , int user ) { if ( locked [ num ] == user ) { locked [ num ] = - 1 ; return true ; } return false ; } bool upgrade ( int num , int user ) { int x = num ; while ( x != - 1 ) { if ( locked [ x ] != - 1 ) { return false ; } x = parent [ x ]; } bool find = false ; function < void ( int ) > dfs = [ & ]( int x ) { for ( int y : children [ x ]) { if ( locked [ y ] != - 1 ) { find = true ; locked [ y ] = - 1 ; } dfs ( y ); } }; dfs ( num ); if ( ! find ) { return false ; } locked [ num ] = user ; return true ; } private: vector < int > locked ; vector < int > parent ; vector < vector < int >> children ; }; /** * Your LockingTree object will be instantiated and called as such: * LockingTree* obj = new LockingTree(parent); * bool param_1 = obj->lock(num,user); * bool param_2 = obj->unlock(num,user); * bool param_3 = obj->upgrade(num,user); */
```

### Python

```python
class LockingTree : def __init__ ( self , parent : List [ int ]): n = len ( parent ) self . locked = [ - 1 ] * n self . parent = parent self . children = [[] for _ in range ( n )] for son , fa in enumerate ( parent [ 1 :], 1 ): self . children [ fa ]. append ( son ) def lock ( self , num : int , user : int ) -> bool : if self . locked [ num ] == - 1 : self . locked [ num ] = user return True return False def unlock ( self , num : int , user : int ) -> bool : if self . locked [ num ] == user : self . locked [ num ] = - 1 return True return False def upgrade ( self , num : int , user : int ) -> bool : 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 ) x = num while x != - 1 : if self . locked [ x ] != - 1 : return False x = self . parent [ x ] find = False dfs ( num ) if not find : return False self . locked [ num ] = user return True # Your LockingTree object will be instantiated and called as such: # obj = LockingTree(parent) # param_1 = obj.lock(num,user) # param_2 = obj.unlock(num,user) # param_3 = obj.upgrade(num,user)
```