Skip to content

1993. Operations on Tree 👍

  • Time:
    • Constructor: $O(n)$
    • lock(num: int, user: int): $O(1)$
    • unlock(num: int, user: int): $O(1)$
    • upgrade(num: int, user: int): $O(n)$
  • Space: $O(n)$
 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
struct Node {
  vector<int> children;
  int lockedBy = -1;
};

class LockingTree {
 public:
  LockingTree(vector<int>& parent) : parent(parent) {
    nodes.resize(parent.size());
    for (int i = 1; i < parent.size(); ++i)
      nodes[parent[i]].children.push_back(i);
  }

  bool lock(int num, int user) {
    if (nodes[num].lockedBy != -1)
      return false;
    return nodes[num].lockedBy = user;
  }

  bool unlock(int num, int user) {
    if (nodes[num].lockedBy != user)
      return false;
    return nodes[num].lockedBy = -1;
  }

  bool upgrade(int num, int user) {
    if (nodes[num].lockedBy != -1)
      return false;
    if (!anyLockedDescendant(num))
      return false;

    // Walk up the hierarchy to ensure that there are no locked ancestors.
    for (int i = num; i != -1; i = parent[i])
      if (nodes[i].lockedBy != -1)
        return false;

    unlockDescendants(num);
    return nodes[num].lockedBy = user;
  }

 private:
  const vector<int> parent;
  vector<Node> nodes;

  bool anyLockedDescendant(int i) {
    return nodes[i].lockedBy != -1 ||
           ranges::any_of(nodes[i].children, [this](const int child) {
      return anyLockedDescendant(child);
    });
  }

  void unlockDescendants(int i) {
    nodes[i].lockedBy = -1;
    for (const int child : nodes[i].children)
      unlockDescendants(child);
  }
};

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
class Node {
  public List<Integer> children = new ArrayList<>();
  public int lockedBy = -1;
}

class LockingTree {
  public LockingTree(int[] parent) {
    this.parent = parent;
    nodes = new Node[parent.length];
    for (int i = 0; i < parent.length; ++i)
      nodes[i] = new Node();
    for (int i = 1; i < parent.length; ++i)
      nodes[parent[i]].children.add(i);
  }

  public boolean lock(int num, int user) {
    if (nodes[num].lockedBy != -1)
      return false;
    nodes[num].lockedBy = user;
    return true;
  }

  public boolean unlock(int num, int user) {
    if (nodes[num].lockedBy != user)
      return false;
    nodes[num].lockedBy = -1;
    return true;
  }

  public boolean upgrade(int num, int user) {
    if (nodes[num].lockedBy != -1)
      return false;
    if (!anyLockedDescendant(num))
      return false;

    // Walk up the hierarchy to ensure that there are no locked ancestors.
    for (int i = num; i != -1; i = parent[i])
      if (nodes[i].lockedBy != -1)
        return false;

    unlockDescendants(num);
    nodes[num].lockedBy = user;
    return true;
  }

  private int[] parent;
  private Node[] nodes;

  private boolean anyLockedDescendant(int i) {
    return nodes[i].lockedBy != -1 ||
        nodes[i].children.stream().anyMatch(child -> anyLockedDescendant(child));
  }

  private void unlockDescendants(int i) {
    nodes[i].lockedBy = -1;
    for (final int child : nodes[i].children)
      unlockDescendants(child);
  }
}

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
class Node:
  def __init__(self):
    self.children: list[int] = []
    self.lockedBy = -1


class LockingTree:
  def __init__(self, parent: list[int]):
    self.parent = parent
    self.nodes = [Node() for _ in range(len(parent))]
    for i in range(1, len(parent)):
      self.nodes[parent[i]].children.append(i)

  def lock(self, num: int, user: int) -> bool:
    if self.nodes[num].lockedBy != -1:
      return False
    self.nodes[num].lockedBy = user
    return True

  def unlock(self, num: int, user: int) -> bool:
    if self.nodes[num].lockedBy != user:
      return False
    self.nodes[num].lockedBy = -1
    return True

  def upgrade(self, num: int, user: int) -> bool:
    if self.nodes[num].lockedBy != -1:
      return False
    if not self._anyLockedDescendant(num):
      return False

    # Walk up the hierarchy to ensure that there are no locked ancestors.
    i = num
    while i != -1:
      if self.nodes[i].lockedBy != -1:
        return False
      i = self.parent[i]

    self._unlockDescendants(num)
    self.nodes[num].lockedBy = user
    return True

  def _anyLockedDescendant(self, i: int) -> bool:
    return (self.nodes[i].lockedBy != -1 or
            any(self._anyLockedDescendant(child)
            for child in self.nodes[i].children))

  def _unlockDescendants(self, i: int) -> None:
    self.nodes[i].lockedBy = -1
    for child in self.nodes[i].children:
      self._unlockDescendants(child)