LeetCode-1993 树上的操作

class LockingTree:

    def __init__(self, parent: List[int]):
        self.n = len(parent)
        self.parent = parent
        self.is_locked = [0 for _ in range(self.n)]
        self.locked_by = [-1 for _ in range(self.n)]
        self.locked_childs = [[] for _ in range(self.n)]
        self.count = 0

    def lock(self, num: int, user: int) -> bool:
        if self.is_locked[num] == 1:
            self.count += 1
            return False
        self.is_locked[num] = 1
        self.locked_by[num] = user   
        x = num 
        while self.parent[x] != -1:
            x = self.parent[x]
            self.locked_childs[x].append(num)
        self.count += 1
        return True

    def unlock(self, num: int, user: int) -> bool:
        if user != self.locked_by[num]:
            self.count += 1
            return False
        if self.is_locked[num] == 0:
            self.count += 1
            return False
        self.is_locked[num] = 0
        x = num
        while self.parent[x] != -1:
            x = self.parent[x]
            self.locked_childs[x].remove(num)
        self.count += 1
        return True

    def upgrade(self, num: int, user: int) -> bool:
        if self.is_locked[num] == 1 or len(self.locked_childs[num]) == 0:
            self.count += 1
            return False
        is_valid = 1
        x = num
        while self.parent[x] != -1:
            x = self.parent[x]
            if self.is_locked[x] == 1:
                is_valid = 0
                self.count += 1
                return False

        self.lock(num, user)
        child_list = self.locked_childs[num][:]
        for child in child_list:
            self.unlock(child, self.locked_by[child])
        self.count += 1
        return True