1
class LockingTree:
2

3
    def __init__(self, parent: List[int]):
4
        self.p = collections.defaultdict(lambda: -2)
5
        self.c = collections.defaultdict(list)
6
        for i, p in enumerate(parent):
7
            self.c[p].append(i)
8
            self.p[i] = p
9
        self.user = collections.defaultdict(set)
10
        self.node = collections.defaultdict(lambda: -2)
11

12
    def lock(self, num: int, user: int) -> bool:
13
        if self.node[num] == -2:
14
            self.user[user].add(num)
15
            self.node[num] = user
16
            return True
17
        return False
18

19
    def unlock(self, num: int, user: int) -> bool:
20
        if self.node[num] == user:
21
            del self.node[num]
22
            self.user[user].remove(num)
23
            return True
24
        return False
25

26
    def upgrade(self, num: int, user: int) -> bool:
27
        if self.node[num] != -2:
28
            return False
29
        if not self.has_locked_descendant(num):
30
            return False
31
        if self.has_locked_ancester(num):
32
            return False
33
        self.lock(num, user)
34
        self.unlock_descendant(num)
35
        return True
36

37
    def has_locked_descendant(
38
        self, num
39
    ):  # function to check if alteast one desendent is lock or not
40
        has = False
41
        for child in self.c[num]:
42
            if self.node[child] != -2:
43
                return True
44
            has |= self.has_locked_descendant(child)
45
        return has
46

47
    def has_locked_ancester(self, num):  # function to check if no parent is locked
48
        if num == -1:
49
            return False
50
        if self.node[self.p[num]] != -2:
51
            return True
52
        return self.has_locked_ancester(self.p[num])
53

54
    def unlock_descendant(self, num):  # function fro unlocking all desendents
55
        for child in self.c[num]:
56
            if child in self.node:
57
                user = self.node[child]
58
                del self.node[child]
59
                if user in self.user:
60
                    self.user[user].remove(child)
61
            self.unlock_descendant(child)