1
class LockingTree {
2
  int[] p;
3
  Map<Integer, Integer> map = new HashMap<>();
4
  Map<Integer, List<Integer>> children = new HashMap<>();
5

6
  public LockingTree(int[] parent) {
7
    p = parent;
8
    for (int i = 0; i < p.length; i++) {
9
      children.put(i, new ArrayList<>());
10
    }
11
    for (int i = 1; i < p.length; i++) {
12
      children.get(p[i]).add(i);
13
    }
14
  }
15

16
  public boolean lock(int num, int user) {
17
    if (!map.containsKey(num)) {
18
      map.put(num, user);
19
      return true;
20
    }
21
    return false;
22
  }
23

24
  public boolean unlock(int num, int user) {
25
    if (map.containsKey(num) && map.get(num) == user) {
26
      map.remove(num);
27
      return true;
28
    }
29
    return false;
30
  }
31

32
  public boolean upgrade(int num, int user) {
33
    // check the node
34
    if (map.containsKey(num)) return false;
35
    // check Ancestor
36
    int ori = num;
37
    while (p[num] != -1) {
38
      if (map.get(p[num]) != null) return false;
39
      num = p[num];
40
    }
41
    // check Decendant
42
    Queue<Integer> q = new LinkedList<>();
43
    List<Integer> child = children.get(ori);
44
    if (child != null) {
45
      for (int c : child) q.offer(c);
46
    }
47
    boolean lock = false;
48
    while (!q.isEmpty()) {
49
      int cur = q.poll();
50
      if (map.get(cur) != null) {
51
        lock = true;
52
        map.remove(cur); // unlock
53
      }
54
      List<Integer> cc = children.get(cur);
55
      if (cc != null) {
56
        for (int c : cc) q.offer(c);
57
      }
58
    }
59
    if (!lock) return false;
60
    map.put(ori, user); // lock the original node
61
    return true;
62
  }
63
}