树上的操作

难度: Medium

给你一棵 n 个节点的树，编号从 0 到 n - 1 ，以父节点数组 parent 的形式给出，其中 parent[i] 是第 i 个节点的父节点。树的根节点为 0 号节点，所以 parent[0] = -1 ，因为它没有父节点。你想要设计一个数据结构实现树里面对节点的加锁，解锁和升级操作。

数据结构需要支持如下函数：

Lock：指定用户给指定节点上锁，上锁后其他用户将无法给同一节点上锁。只有当节点处于未上锁的状态下，才能进行上锁操作。
Unlock：指定用户给指定节点解锁，只有当指定节点当前正被指定用户锁住时，才能执行该解锁操作。
Upgrade：指定用户给指定节点上锁，并且将该节点的所有子孙节点解锁。只有如下 3 个条件全部满足时才能执行升级操作：
- 指定节点当前状态为未上锁。
- 指定节点至少有一个上锁状态的子孙节点（可以是任意用户上锁的）。
- 指定节点没有任何上锁的祖先节点。

请你实现 LockingTree 类：

LockingTree(int[] parent) 用父节点数组初始化数据结构。
lock(int num, int user) 如果 id 为 user 的用户可以给节点 num 上锁，那么返回 true ，否则返回 false 。如果可以执行此操作，节点 num 会被 id 为 user 的用户上锁。
unlock(int num, int user) 如果 id 为 user 的用户可以给节点 num 解锁，那么返回 true ，否则返回 false 。如果可以执行此操作，节点 num 变为 未上锁 状态。
upgrade(int num, int user) 如果 id 为 user 的用户可以给节点 num 升级，那么返回 true ，否则返回 false 。如果可以执行此操作，节点 num 会被升级。

示例 1：

输入：
["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]]
输出：
[null, true, false, true, true, true, false]

解释：
LockingTree lockingTree = new LockingTree([-1, 0, 0, 1, 1, 2, 2]);
lockingTree.lock(2, 2);    // 返回 true ，因为节点 2 未上锁。
                           // 节点 2 被用户 2 上锁。
lockingTree.unlock(2, 3);  // 返回 false ，因为用户 3 无法解锁被用户 2 上锁的节点。
lockingTree.unlock(2, 2);  // 返回 true ，因为节点 2 之前被用户 2 上锁。
                           // 节点 2 现在变为未上锁状态。
lockingTree.lock(4, 5);    // 返回 true ，因为节点 4 未上锁。
                           // 节点 4 被用户 5 上锁。
lockingTree.upgrade(0, 1); // 返回 true ，因为节点 0 未上锁且至少有一个被上锁的子孙节点（节点 4）。
                           // 节点 0 被用户 1 上锁，节点 4 变为未上锁。
lockingTree.lock(0, 1);    // 返回 false ，因为节点 0 已经被上锁了。

提示：

n == parent.length
2 <= n <= 2000
对于 i != 0 ，满足 0 <= parent[i] <= n - 1
parent[0] == -1
0 <= num <= n - 1
1 <= user <= 10⁴
parent 表示一棵合法的树。
lock ，unlock 和 upgrade 的调用总共不超过 2000 次。

Submission

运行时间: 416 ms

内存: 18.7 MB

class LockingTree:

    def __init__(self, parent: List[int]):
        n=len(parent)
        self.parent=parent
        self.g=[[] for _ in range(n)]
        for x,y in enumerate(parent):
            if y!=-1:
                self.g[y].append(x)
        self.locked=[False]*n
        self.users=[-1]*n
        self.control=[0]*n##统计存在上锁节点的子树数量

    def setLock(self,num: int, user: int):
        self.locked[num]=True
        self.users[num]=user

    def lock(self, num: int, user: int) -> bool:
        if self.locked[num]:return False
        self.setLock(num,user)
        if self.control[num]==0:
            node=self.parent[num]
            while node!=-1:
                self.control[node]+=1
                if self.locked[node] or self.control[node]>1:
                    break
                node=self.parent[node]
        return True

    def unlock(self, num: int, user: int) -> bool:
        if not self.locked[num] or self.users[num]!=user:
            return False
        self.locked[num]=False
        if self.control[num]==0:
            node=self.parent[num]
            while node!=-1:
                self.control[node]-=1
                if self.locked[node] or self.control[node]>=1:
                    break
                node=self.parent[node]
        return True

    def upgrade(self, num: int, user: int) -> bool:
        def subNodeUnlock(node):
            for y in self.g[node]:
                self.locked[y]=False
                if self.control[y]>0:
                    self.control[y]=0
                    subNodeUnlock(y)
        
        if self.locked[num] or self.control[num]==0:
            return False
        cur=self.parent[num]
        while cur!=-1:
            if self.locked[cur]:return False
            cur=self.parent[cur]
        self.setLock(num,user)
        if self.control[num]>0:
            self.control[num]=0
            subNodeUnlock(num)
        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)

Explain

该数据结构设计的核心在于管理一棵树的节点锁定状态。`LockingTree` 类使用数组 `locked` 来表示每个节点是否被锁定，使用 `users` 记录锁定每个节点的用户ID。此外，还使用了一个 `control` 数组来记录每个节点的子树中被锁定的节点数量。树的结构由 `parent` 数组和 `g` 列表（邻接表）表示。锁定 (`lock`) 和解锁 (`unlock`) 操作都会更新 `control` 数组，以便快速判断节点的锁定状态。`upgrade` 操作则需要检查目标节点自身未锁定、无锁定的祖先节点和至少一个锁定的子孙节点这三个条件，满足条件后将该节点锁定并解锁所有子孙节点。

时间复杂度: O(n)

空间复杂度: O(n)

class LockingTree:

    def __init__(self, parent: List[int]):
        n=len(parent)
        self.parent=parent
        self.g=[[] for _ in range(n)]  # 邻接表表示树
        for x,y in enumerate(parent):
            if y!=-1:
                self.g[y].append(x)
        self.locked=[False]*n  # 节点锁定状态
        self.users=[-1]*n  # 锁定节点的用户
        self.control=[0]*n  # 子树中锁定节点的数量

    def setLock(self,num: int, user: int):
        self.locked[num]=True
        self.users[num]=user

    def lock(self, num: int, user: int) -> bool:
        if self.locked[num]:return False
        self.setLock(num,user)
        if self.control[num]==0:
            node=self.parent[num]
            while node!=-1:
                self.control[node]+=1
                if self.locked[node] or self.control[node]>1:
                    break
                node=self.parent[node]
        return True

    def unlock(self, num: int, user: int) -> bool:
        if not self.locked[num] or self.users[num]!=user:
            return False
        self.locked[num]=False
        if self.control[num]==0:
            node=self.parent[num]
            while node!=-1:
                self.control[node]-=1
                if self.locked[node] or self.control[node]>=1:
                    break
                node=self.parent[node]
        return True

    def upgrade(self, num: int, user: int) -> bool:
        def subNodeUnlock(node):
            for y in self.g[node]:
                self.locked[y]=False
                if self.control[y]>0:
                    self.control[y]=0
                    subNodeUnlock(y)
        
        if self.locked[num] or self.control[num]==0:
            return False
        cur=self.parent[num]
        while cur!=-1:
            if self.locked[cur]:return False
            cur=self.parent[cur]
        self.setLock(num,user)
        if self.control[num]>0:
            self.control[num]=0
            subNodeUnlock(num)
        return True

Explore

在`LockingTree`类的构造函数中，邻接表`g`通过遍历`parent`数组来构建。对于`parent`数组中的每一对索引和值（即每一个节点及其父节点），如果父节点不是-1（表示该节点不是根节点），则将该节点添加到其父节点在`g`列表对应的子列表中。这样构建的邻接表`g`反映了树的结构，其中每个节点列表存储了其所有直接子节点。在整个类中，`g`主要用于表示树的结构，特别是在执行操作如`upgrade`时，递归访问某节点的所有子孙节点以检查或修改它们的锁定状态时使用。

在`lock`方法中，更新`control`数组是为了记录从当前节点向上到根节点路径上，各个节点的子树中锁定节点的数量。检查`if self.locked[node] or self.control[node]>1`的条件用来确定何时停止更新`control`数组。具体来说，如果某个祖先节点已经被锁定或者其`control`值大于1，意味着在此之上的祖先节点的`control`计数已经包含了当前节点的锁定状态，因此无需继续更新。这防止了重复计数并且减少了不必要的更新，优化了操作的效率。

在`unlock`方法中，检查`if not self.locked[num] or self.users[num]!=user`是为了确保只有锁定了特定节点的用户才能解锁它。首先，检查`self.locked[num]`确认该节点当前是否被锁定；其次，通过比较`self.users[num]`和`user`确认尝试执行解锁操作的用户是否是锁定该节点的用户。这样的检查防止了未经授权的用户解锁不属于他们锁定的节点，确保了操作的安全性和正确性。