Tree of Space

Using Hashmap

class Node:
    name: str
    parent: None
    children: list["Node"]
    isLocked: bool
    lockedBy: int
    numLockedDescendants: int
    def __init__(self, name):
        self.isLocked:bool = False
        self.lockedBy = -1
        self.numLockedDescendants = 0
        self.parent = None
        self.children = []
        self.name = name
 
 
def lock(v, uid):
    if v.isLocked or v.numLockedDescendants > 0:
        return False
 
    parent = v.parent
    while parent != None:
        if parent.isLocked:
            return False
        parent = parent.parent
 
    v.isLocked = True
    v.lockedBy = uid
 
    parent = v.parent
    while parent != None:
        parent.numLockedDescendants += 1
        parent = parent.parent
 
    return True
 
def unlock(v, uid):
    if v.isLocked or not v.lockedBy == uid:
        return False
 
    v.isLocked = False
    v.lockedBy = -1
 
    parent = v.parent
    while parent != None:
        parent.numLockedDescendants -= 1
        parent = parent.parent
 
    return True
 
def collectLockedDescendants(v, uid, lockedList):
    if v.numLockedDescendants == 0 and not v.isLocked:
        return True
 
    if v.isLocked:
        if v.lockedBy != uid:
            return False
        lockedList.append(v)
 
    for child in v.children:
        if not collectLockedDescendants(child, uid, lockedList):
            return False
 
    return True
 
def upgrade(v, uid):
    if v.isLocked and v.numLockedDescendants == 0:
        return False
 
    lockedList = []
    if (collectLockedDescendants(v, uid, lockedList) == False):
        return False
 
    for node in lockedList:
        unlock(node, uid)
 
    v.isLocked = True
    v.lockedBy = uid
 
    parent = v.parent
    while parent != None:
        parent.numLockedDescendants += 1
        parent = parent.parent
 
    return True
 
def build_tree():
    A = Node('A')
    B = Node('B'); C = Node('C')
    D = Node('D'); E = Node('E')
    A.children = [B, C]; B.parent = A; C.parent = A
    B.children = [D, E]; D.parent = B; E.parent = B
    return A, B, C, D, E
 
A, B, C, D, E = build_tree()
 
print(lock(D, 1))      # True — lock leaf D by user 1
print(lock(E, 1))      # True — lock E by user 1
print(upgrade(B, 1))   # True — upgrade locks D & E -> B
print(A.numLockedDescendants)  # Expect 1 (only B locked)
print(lock(D, 2))      # False — ancestor B locked by user 1
print(unlock(B, 1))    # True — unlock B
print(lock(D, 2))      # True — now user 2 can lock D
print(unlock(E, 1))    # False — E isn't locked
 
# Print statuses for verification
for node in [A, B, C, D, E]:
    print(f"{node.name}: locked={node.isLocked}, by={node.lockedBy}, lockedDesc={node.numLockedDescendants}")
 

Leetcode 1993

class LockingTree:
 
  
 
def __init__(self, parent: List[int]):
 
self.parent = parent
 
self.locked = [None] * len(parent)
 
self.child = {i: [] for i in range(len(parent))}
 
for i in range(1, len(parent)):
 
self.child[parent[i]].append(i)
 
  
 
def lock(self, num: int, user: int) -> bool:
 
if self.locked[num]: return False
 
self.locked[num] = user
 
return True
 
  
 
def unlock(self, num: int, user: int) -> bool:
 
if self.locked[num] != user:
 
return False
 
  
 
self.locked[num] = None
 
return True
 
  
 
def upgrade(self, num: int, user: int) -> bool:
 
# Check if there is any locked ancestors
 
temp = num
 
while temp != -1:
 
if self.locked[temp]:
 
return False
 
temp = self.parent[temp]
 
  
 
# Check if atleast one descendant is locked
 
from collections import deque
 
  
 
q = deque()
 
q.append(num)
 
locked = 0
 
while q:
 
val = q.popleft()
 
if self.locked[val]:
 
self.locked[val] = None
 
locked += 1
 
child = self.child[val]
 
q.extend(child)
 
  
 
if locked == 0:
 
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)

Binary Lifting