Problem

Implement locking in a binary tree. A binary tree node can be locked or unlocked only if all of its descendants or ancestors are not locked.

Design a binary tree node class with the following methods:

  • is_locked, which returns whether the node is locked
  • lock, which attempts to lock the node. If it cannot be locked, then it should return false. Otherwise, it should lock it and return true.
  • unlock, which unlocks the node. If it cannot be unlocked, then it should return false. Otherwise, it should unlock it and return true.

You may augment the node to add parent pointers or any other property you would like. You may assume the class is used in a single-threaded program, so there is no need for actual locks or mutexes. Each method should run in O(h), where h is the height of the tree.

Examples

Example 1:

	  1
	/   \
   2     3
  / \   / \
 *   * 4   *

Input: root = [1,2,3, #,#,4,#], 
operations = [ node3.lock()], node3.is_locked(), root.lock(), node4.lock(), node3.unlock(), node4.lock()]
Output: [true, true, false, false, true, true]
Explanation:
node3.lock() outputs true because lock was successful
node3.is_locked() outputs true because it is locked
root.lock() outputs false as one of the descendants, i.e. node 3 is locked
node4.lock() outputs false because one of the ascendants was locked i.e. node 3
node3.unlock() outputs true as unlocking was successful
nod4.lock() outputs true because lock was successful

Solution

Method 1 - Preorder Traversal

Code

Java
class TreeNode {
	int val;
	boolean isLocked;
	boolean hasLockedDescendant; // Tracks descendants' locking state
	TreeNode left;
	TreeNode right;

	public TreeNode(int val, boolean isLocked) {
		this.val = val;
		this.isLocked = isLocked;
		this.hasLockedDescendant = false;
	}

	public boolean isLocked() {
		return isLocked;
	}

	public boolean lock() {
		if (isLocked || hasLockedDescendant) {
			return false;
		}

		isLocked = true;
		updateHasLockedDescendant(true); // Propagate lock state down
		return true;
	}

	public boolean unlock() {
		if (!isLocked) {
			return false;
		}

		isLocked = false;
		updateHasLockedDescendant(false); // Propagate unlock state up
		return true;
	}

	private void updateHasLockedDescendant(boolean locked) {
		hasLockedDescendant = locked;

		if (left != null) {
			left.updateHasLockedDescendant(locked);
		}

		if (right != null) {
			right.updateHasLockedDescendant(locked);
		}
	}
}
Python
class TreeNode:
  def __init__(self, is_locked):
    self.is_locked = is_locked
    self.has_locked_descendant = False
    self.left = None
    self.right = None

  def is_locked(self):
    return self.is_locked

  def lock(self):
    if self.is_locked or self.has_locked_descendant:
      return False
    self.is_locked = True
    self.update_has_locked_descendant(True)
    return True

  def unlock(self):
    if not self.is_locked:
      return False
    self.is_locked = False
    self.update_has_locked_descendant(False)
    return True

  def update_has_locked_descendant(self, locked):
    self.has_locked_descendant = locked
    if self.left:
      self.left.update_has_locked_descendant(locked)
    if self.right:
      self.right.update_has_locked_descendant(locked)

Complexity

  • ⏰ Time complexity: O(h) where h is the height of the tree. This is because they need to traverse down or up the tree at most h levels in the worst case to update the locking state of descendants or ancestors.
  • 🧺 Space complexity: O(h), assuming recursion stack.