Problem statement

A tree of n nodes, numbered 0 to n - 1, is given as a parent array: parent[i] is the parent of node i, and the root 0 has parent[0] == -1. Users lock and unlock its nodes under these rules:

• Lock — a node can be locked by a user only if it is currently unlocked. It then belongs to that user.
• Unlock — a node can be unlocked only by the user who locked it.
• Upgrade — a user may lock a node and at the same time unlock all of its descendants, but only if all three hold: the node is unlocked, it has at least one locked descendant (locked by anyone), and none of its ancestors is locked.

Write the class LockingTree:

• LockingTree(int[] parent) builds the tree, all nodes unlocked.
• boolean lock(int num, int user) locks node num for user if allowed; returns whether it did.
• boolean unlock(int num, int user) unlocks node num if user holds it; returns whether it did.
• boolean upgrade(int num, int user) performs the upgrade on num for user if allowed; returns whether it did.

The judge calls the methods in the order listed in the first input line, with the arguments in the second; the expected output holds each call's result, null for the constructor and for every method that returns nothing.

Example 1

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: Node 2 is locked by user 2; user 3 cannot unlock it, user 2 can. After user 5 locks node 4, user 1 upgrades the root: it is unlocked, has the locked descendant 4 and no ancestors, so node 4 is unlocked and the root locked for user 1. Locking the root again then fails.

Example 2

Input: ["LockingTree","lock","upgrade","upgrade","lock","upgrade"]
[[[-1,0,1,2]],[3,7],[2,4],[0,9],[1,5],[2,6]]
Output: [null,true,true,true,true,false]
Explanation: The tree is the chain 0-1-2-3 and node 3 is locked. Upgrading 2 succeeds and unlocks 3. Upgrading 0 then succeeds too — 2 is its locked descendant, and is unlocked by it. Once node 1 is locked, upgrading 2 fails: its ancestors 1 and 0 are locked.

Constraints

• n == parent.length
• 2 <= n <= 2000
• 0 <= parent[i] <= n - 1 for i != 0, and parent[0] == -1
• parent describes a valid tree
• 0 <= num <= n - 1
• 1 <= user <= 10^4
• At most 2000 calls are made to lock, unlock and upgrade in total

Hints for Operations on Tree

1. Hint 1

   Keep, for each node, the user who holds it (0 for unlocked). Lock and unlock are then single checks.

2. Hint 2

   For upgrade, the ancestor condition is a walk up the parent pointers. The descendant condition needs to walk down — build each node's list of children once in the constructor.

3. Hint 3

   Check the node and its ancestors first. Then traverse the subtree below num (DFS or BFS), unlocking every locked node you meet and noting whether there was one; if there was none, nothing changed and the answer is false, otherwise lock num for the user.

How to solve Operations on Tree

Store the locking user per node and each node's children. Upgrade is a walk up (ancestors) followed by a traversal down (descendants), both bounded by the tree's size.

Approach

1. Constructor: lockedBy[v] = 0 for every node; children[parent[i]].append(i) for i >= 1.
2. lock(num, user): if lockedBy[num] != 0 return false; else set it to user and return true.
3. unlock(num, user): if lockedBy[num] != user return false; else set it to 0 and return true.
4. upgrade(num, user): return false if num is locked or any node on the path parent[num], parent[parent[num]], … is locked.
5. Traverse the subtree under num with a stack: every locked node found is unlocked and counted. If the count is 0 return false; otherwise set lockedBy[num] = user and return true.

Why it works

The three upgrade conditions are checked exactly as stated: the first two reads leave the tree unchanged when they fail, and the traversal unlocks descendants only — which is harmless when none are locked, since then nothing is modified before the false is returned. Unlocking during the same traversal that searches for a locked descendant saves a second pass.

Complexity

• Time — O(1) for lock and unlock, O(n) for upgrade
• Space — O(n)

Pitfalls

• Unlock must compare against the user who locked the node; any other user fails even if the node is locked.
• Upgrade needs a locked descendant; a locked node elsewhere in the tree does not count, and neither does num itself.
• Check the ancestors before unlocking anything — otherwise a failing upgrade would already have changed descendants.

Reference solution

Python

class LockingTree:

    def __init__(self, parent: List[int]):
        self.parent = parent
        self.children = [[] for _ in parent]
        for v in range(1, len(parent)):
            self.children[parent[v]].append(v)
        self.locked_by = [0] * len(parent)  # 0 = unlocked

    def lock(self, num: int, user: int) -> bool:
        if self.locked_by[num]:
            return False
        self.locked_by[num] = user
        return True

    def unlock(self, num: int, user: int) -> bool:
        if self.locked_by[num] != user:
            return False
        self.locked_by[num] = 0
        return True

    def upgrade(self, num: int, user: int) -> bool:
        lb = self.locked_by
        if lb[num]:
            return False
        a = self.parent[num]
        while a != -1:  # no locked ancestor
            if lb[a]:
                return False
            a = self.parent[a]
        found = False
        stack = list(self.children[num])
        while stack:  # unlock every locked descendant
            v = stack.pop()
            if lb[v]:
                lb[v] = 0
                found = True
            stack.extend(self.children[v])
        if not found:
            return False
        lb[num] = user
        return True

Java

class LockingTree {
    private int[] parent, lockedBy; // lockedBy 0 = unlocked
    private List<List<Integer>> children = new ArrayList<>();

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

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

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

    public boolean upgrade(int num, int user) {
        if (lockedBy[num] != 0) return false;
        for (int a = parent[num]; a != -1; a = parent[a]) if (lockedBy[a] != 0) return false; // no locked ancestor
        boolean found = false;
        Deque<Integer> stack = new ArrayDeque<>(children.get(num));
        while (!stack.isEmpty()) { // unlock every locked descendant
            int v = stack.pop();
            if (lockedBy[v] != 0) { lockedBy[v] = 0; found = true; }
            for (int c : children.get(v)) stack.push(c);
        }
        if (!found) return false;
        lockedBy[num] = user;
        return true;
    }
}

C++

class LockingTree {
    vector<int> parent, lockedBy; // lockedBy 0 = unlocked
    vector<vector<int>> children;
public:
    LockingTree(vector<int>& parent) : parent(parent), lockedBy(parent.size(), 0), children(parent.size()) {
        for (int v = 1; v < (int)parent.size(); v++) children[parent[v]].push_back(v);
    }

    bool lock(int num, int user) {
        if (lockedBy[num]) return false;
        lockedBy[num] = user;
        return true;
    }

    bool unlock(int num, int user) {
        if (lockedBy[num] != user) return false;
        lockedBy[num] = 0;
        return true;
    }

    bool upgrade(int num, int user) {
        if (lockedBy[num]) return false;
        for (int a = parent[num]; a != -1; a = parent[a]) if (lockedBy[a]) return false; // no locked ancestor
        bool found = false;
        vector<int> stack(children[num].begin(), children[num].end());
        while (!stack.empty()) { // unlock every locked descendant
            int v = stack.back();
            stack.pop_back();
            if (lockedBy[v]) { lockedBy[v] = 0; found = true; }
            for (int c : children[v]) stack.push_back(c);
        }
        if (!found) return false;
        lockedBy[num] = user;
        return true;
    }
};

JavaScript

var LockingTree = function(parent) {
    this.parent = parent;
    this.children = parent.map(function() { return []; });
    for (var v = 1; v < parent.length; v++) this.children[parent[v]].push(v);
    this.lockedBy = new Array(parent.length).fill(0); // 0 = unlocked
};

LockingTree.prototype.lock = function(num, user) {
    if (this.lockedBy[num]) return false;
    this.lockedBy[num] = user;
    return true;
};

LockingTree.prototype.unlock = function(num, user) {
    if (this.lockedBy[num] !== user) return false;
    this.lockedBy[num] = 0;
    return true;
};

LockingTree.prototype.upgrade = function(num, user) {
    var lb = this.lockedBy;
    if (lb[num]) return false;
    for (var a = this.parent[num]; a !== -1; a = this.parent[a]) if (lb[a]) return false; // no locked ancestor
    var found = false, stack = this.children[num].slice();
    while (stack.length) { // unlock every locked descendant
        var v = stack.pop();
        if (lb[v]) { lb[v] = 0; found = true; }
        for (var i = 0; i < this.children[v].length; i++) stack.push(this.children[v][i]);
    }
    if (!found) return false;
    lb[num] = user;
    return true;
};

Also on the editorial tab: C, C#, Go, Kotlin, PHP, Ruby, Rust, Swift, TypeScript.

More arrays problems

• Implement Router Medium
• Design Memory Allocator Medium
• Design a Food Rating System Medium
• Design an ATM Machine Medium
• Design Bitset Medium
• Range Frequency Queries Medium