1
class LockingTree {
2
 public:
3
  unordered_map<int, vector<int>> descendents;
4
  vector<vector<int>> Node;
5
  /*
6
      Node[i][0] = parent[i]
7
      Node[i][1] = -1; (means unlocked)
8
      Node[i][1] = x;  (means locked by user x)
9
  */
10
  int n;
11
  LockingTree(vector<int> &parent) {
12
    n = parent.size();
13
    Node.resize(n, vector<int>(2, -1));
14

15
    Node[0][0] = -1; // root has no parent
16
    for (int i = 1; i < n; i++) {
17
      Node[i][0] = parent[i];
18
      descendents[parent[i]].push_back(i);
19
    }
20
  }
21

22
  bool lock(int num, int user) {
23
    if (Node[num][1] != -1) return false;
24

25
    Node[num][1] = user;
26
    return true;
27
  }
28

29
  bool unlock(int num, int user) {
30
    if (Node[num][1] != user) return false;
31

32
    Node[num][1] = -1;
33
    return true;
34
  }
35

36
  // Condition-2 (Atleast one descendent should be locked)
37
  void checkDescendents(int num, bool &atleastOne) {
38
    if (descendents.count(num) == 0 || descendents[num].size() == 0) return;
39

40
    for (int &x : descendents[num]) {
41
      if (Node[x][1] != -1) {
42
        atleastOne = true;
43
        return;
44
      }
45
      checkDescendents(x, atleastOne);
46
    }
47
  }
48

49
  // Condition-3 (Check if any ancestor is locked)
50
  bool IsAnyAncestorLocked(int &num) {
51
    if (num == -1) return false; // you reached end and found none locked
52

53
    return Node[num][1] != -1 || IsAnyAncestorLocked(Node[num][0]);
54
  }
55

56
  void unlockDescendents(int num) {
57
    if (descendents.count(num) == 0 || descendents[num].size() == 0) return;
58

59
    for (int &x : descendents[num]) {
60
      Node[x][1] = -1;
61
      unlockDescendents(x);
62
    }
63
  }
64

65
  bool upgrade(int num, int user) {
66
    // condition : 1
67
    if (Node[num][1] != -1) return false;
68

69
    // condition : 2
70
    bool atleastOne = false;
71
    checkDescendents(num, atleastOne);
72
    // If no node was locked, return false
73
    if (!atleastOne) return false;
74

75
    // condition : 3
76
    if (IsAnyAncestorLocked(Node[num][0])) return false;
77

78
    // Do the rest
79
    unlockDescendents(num);
80
    Node[num][1] = user;
81
    return true;
82
  }
83
};