Operations on Tree
Problem
You are given a rooted tree of n nodes (0 to n−1) described by a parent array, where parent[i] is the parent of node i and the root has parent −1. Design a structure with three operations. lock(num, user) locks node num for user only if it is currently unlocked. unlock(num, user) unlocks node num only if it is currently locked by that same user. upgrade(num, user) locks node num for user only if node num is unlocked, has at least one locked descendant, and has no locked ancestor; on success it also unlocks every descendant. Each call returns true on success and false otherwise.
parent = [-1,0,0,1,1,2,2]; lock(2,2); unlock(2,3); unlock(2,2); lock(4,5); upgrade(0,1)[true, false, true, true, true]class LockingTree:
def __init__(self, parent):
self.parent = parent
self.lock = [0] * len(parent)
self.children = [[] for _ in parent]
for i, p in enumerate(parent):
if p != -1:
self.children[p].append(i)
def lock(self, num, user):
if self.lock[num]:
return False
self.lock[num] = user
return True
def unlock(self, num, user):
if self.lock[num] != user:
return False
self.lock[num] = 0
return True
def upgrade(self, num, user):
a = self.parent[num]
while a != -1:
if self.lock[a]:
return False
a = self.parent[a]
found = False
stack = [num]
while stack:
node = stack.pop()
for c in self.children[node]:
if self.lock[c]:
found = True
self.lock[c] = 0
stack.append(c)
if self.lock[num] or not found:
return False
self.lock[num] = user
return Trueclass LockingTree {
constructor(parent) {
this.parent = parent;
this.lock = new Array(parent.length).fill(0);
this.children = parent.map(() => []);
parent.forEach((p, i) => { if (p !== -1) this.children[p].push(i); });
}
lock(num, user) {
if (this.lock[num]) return false;
this.lock[num] = user;
return true;
}
unlock(num, user) {
if (this.lock[num] !== user) return false;
this.lock[num] = 0;
return true;
}
upgrade(num, user) {
let a = this.parent[num];
while (a !== -1) {
if (this.lock[a]) return false;
a = this.parent[a];
}
let found = false;
const stack = [num];
while (stack.length) {
const node = stack.pop();
for (const c of this.children[node]) {
if (this.lock[c]) { found = true; this.lock[c] = 0; }
stack.push(c);
}
}
if (this.lock[num] || !found) return false;
this.lock[num] = user;
return true;
}
}class LockingTree {
int[] parent, lock;
List<List<Integer>> children;
public LockingTree(int[] parent) {
this.parent = parent;
lock = new int[parent.length];
children = new ArrayList<>();
for (int i = 0; i < parent.length; i++) children.add(new ArrayList<>());
for (int i = 0; i < parent.length; i++)
if (parent[i] != -1) children.get(parent[i]).add(i);
}
public boolean lock(int num, int user) {
if (lock[num] != 0) return false;
lock[num] = user;
return true;
}
public boolean unlock(int num, int user) {
if (lock[num] != user) return false;
lock[num] = 0;
return true;
}
public boolean upgrade(int num, int user) {
for (int a = parent[num]; a != -1; a = parent[a])
if (lock[a] != 0) return false;
boolean found = false;
Deque<Integer> stack = new ArrayDeque<>();
stack.push(num);
while (!stack.isEmpty()) {
int node = stack.pop();
for (int c : children.get(node)) {
if (lock[c] != 0) { found = true; lock[c] = 0; }
stack.push(c);
}
}
if (lock[num] != 0 || !found) return false;
lock[num] = user;
return true;
}
}class LockingTree {
vector<int> parent, lock_;
vector<vector<int>> children;
public:
LockingTree(vector<int>& parent) : parent(parent) {
lock_.assign(parent.size(), 0);
children.resize(parent.size());
for (int i = 0; i < (int)parent.size(); i++)
if (parent[i] != -1) children[parent[i]].push_back(i);
}
bool lock(int num, int user) {
if (lock_[num]) return false;
lock_[num] = user;
return true;
}
bool unlock(int num, int user) {
if (lock_[num] != user) return false;
lock_[num] = 0;
return true;
}
bool upgrade(int num, int user) {
for (int a = parent[num]; a != -1; a = parent[a])
if (lock_[a]) return false;
bool found = false;
vector<int> stack = { num };
while (!stack.empty()) {
int node = stack.back(); stack.pop_back();
for (int c : children[node]) {
if (lock_[c]) { found = true; lock_[c] = 0; }
stack.push_back(c);
}
}
if (lock_[num] || !found) return false;
lock_[num] = user;
return true;
}
};Explanation
The whole structure boils down to one array: lock[i] holds the id of the user who locked node i, or 0 when the node is free. (User ids are positive, so 0 is a safe "nobody" sentinel.) From the parent array we also precompute a children list so we can walk downward later.
lock and unlock are then trivial. To lock, the node must be free; we just stamp the user's id on it. To unlock, the node must currently be held by the same user asking — otherwise we refuse and return false.
upgrade is the interesting one and has three conditions, all of which must hold. First, no ancestor is locked: we climb from num up through parent pointers, and if we hit any locked node we bail. Second, the node itself must be free. Third, there must be at least one locked descendant: we run a downward traversal (DFS with a stack here, BFS works just as well) over the subtree, and as we go we both detect a locked descendant and clear every lock we pass.
If all three checks pass, the descendants are already unlocked from the sweep, and we finish by locking num for the user.
Example: with 4 locked and everything else free, upgrade(0,1) first climbs from 0 — it is the root, so there is no locked ancestor. Then the downward sweep finds node 4 locked, sets found = true, and unlocks it. Node 0 is free and a descendant was locked, so the upgrade succeeds and 0 becomes locked for user 1.