Learn AI Create the Future — with Pranay
Main site

LeetCode solutions

2104. Operations on Tree

My accepted Python solution to LeetCode problem 2104, Operations on Tree, running in 544ms.

  • Difficulty: Medium
  • Python
  • Runtime 544ms
  • Memory 21.1MB

Read the problem on LeetCode View on GitHub

The problem statement is LeetCode’s and stays on their site. What follows is my accepted solution.

Python

Accepted on LeetCode — runtime 544ms, memory 21.1MB, accepted 2025-12-29.

python
class LockingTree:

    def __init__(self, parent: List[int]):
        self.parent = parent
        self.n = len(parent)
        self.locked = [-1] * self.n  # -1 means unlocked, otherwise user who locked
        # Build children list
        self.children = [[] for _ in range(self.n)]
        for i in range(1, self.n):
            self.children[parent[i]].append(i)

    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:
        # Check if node is unlocked
        if self.locked[num] != -1:
            return False
        
        # Check if any ancestor is locked
        curr = num
        while curr != -1:
            if self.locked[curr] != -1:
                return False
            curr = self.parent[curr]
        
        # Check if at least one descendant is locked and unlock all
        def unlock_descendants(node):
            count = 0
            if self.locked[node] != -1:
                self.locked[node] = -1
                count = 1
            for child in self.children[node]:
                count += unlock_descendants(child)
            return count
        
        locked_count = unlock_descendants(num)
        if locked_count == 0:
            return False
        
        # Lock the node
        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)

Source