🔐 Tree of Space — Complete Solution (From Scratch)

Both Part A (Lock/Unlock) and Part B (UpgradeLock + Thread Safety discussion) — fully implemented in C++.

Part A Solution — Lock, Unlock, UpgradeLock (Single-Threaded)

Core Design Philosophy

The critical insight is maintaining a lockedDescendants set in every node.

Without it: Every lock() call requires O(N) BFS to find any locked descendant.
With it: lockedDescendants.size() == 0 is O(1).
When a node is locked, it walks up to root (O(log_M N)) and adds itself to each ancestor's lockedDescendants. Unlock does the reverse.

Tree:
          World
        /       \
     Asia       Africa
    /    \     /     \
 China  India  SA   Egypt

When China is locked:
  Asia.lockedDescendants = {China}
  World.lockedDescendants = {China}
  
When India is also locked:
  Asia.lockedDescendants = {China, India}
  World.lockedDescendants = {China, India}

Complete C++ Implementation

#include <bits/stdc++.h>
using namespace std;

// ─── Node Structure ───────────────────────────────────────────────────────────
struct TreeNode {
    string name;
    int lockedBy;                     // -1 = not locked
    bool isLocked;
    vector<TreeNode*> children;
    TreeNode* parent;
    set<TreeNode*> lockedDescendants; // KEY: tracks all locked nodes below

    TreeNode(string nm, TreeNode* par) {
        name     = nm;
        lockedBy = -1;
        isLocked = false;
        parent   = par;
    }
};

// ─── M-Ary Tree ───────────────────────────────────────────────────────────────
struct MAryTree {
    TreeNode* root;
    unordered_map<string, TreeNode*> nodeMap; // O(1) name-to-pointer lookup

    MAryTree(const string& rootName) {
        root = new TreeNode(rootName, nullptr);
    }

    // Build tree in BFS order from input array
    void build(vector<string>& names, int M) {
        int idx = 1;
        int n   = names.size();
        queue<TreeNode*> q;
        q.push(root);

        while (!q.empty()) {
            TreeNode* cur = q.front(); q.pop();
            nodeMap[cur->name] = cur;

            // Assign up to M children
            for (int i = 0; i < M && idx < n; i++, idx++) {
                TreeNode* child = new TreeNode(names[idx], cur);
                cur->children.push_back(child);
                q.push(child);
            }
        }
    }

    // Walk from `start` to root, inserting `node` into each ancestor's set
    // Time: O(log_M N)
    void propagateUp(TreeNode* start, TreeNode* node, bool insert) {
        TreeNode* cur = start;
        while (cur) {
            if (insert) cur->lockedDescendants.insert(node);
            else        cur->lockedDescendants.erase(node);
            cur = cur->parent;
        }
    }

    // ── LOCK ─────────────────────────────────────────────────────────────────
    // Returns true if node X is successfully locked for uid.
    // Conditions to FAIL:
    //   1. X is already locked
    //   2. X has any locked descendants (lockedDescendants non-empty)
    //   3. Any ancestor of X is locked
    bool lock(const string& name, int uid) {
        TreeNode* node = nodeMap[name];

        // Check 1: already locked?
        if (node->isLocked) return false;

        // Check 2: any locked descendants?
        if (!node->lockedDescendants.empty()) return false;

        // Check 3: any locked ancestor?
        TreeNode* anc = node->parent;
        while (anc) {
            if (anc->isLocked) return false;
            anc = anc->parent;
        }

        // All checks passed — perform lock
        node->isLocked = true;
        node->lockedBy = uid;
        propagateUp(node->parent, node, true); // notify ancestors
        return true;
    }

    // ── UNLOCK ───────────────────────────────────────────────────────────────
    // Returns true if node X (locked by uid) is successfully unlocked.
    // Conditions to FAIL:
    //   1. X is not locked
    //   2. X is locked by a different uid
    bool unlock(const string& name, int uid) {
        TreeNode* node = nodeMap[name];

        // Check 1 & 2
        if (!node->isLocked || node->lockedBy != uid) return false;

        // Perform unlock
        node->isLocked = false;
        node->lockedBy = -1;
        propagateUp(node->parent, node, false); // remove from ancestors
        return true;
    }

    // ── UPGRADE LOCK ─────────────────────────────────────────────────────────
    // Returns true if node X can be upgraded: all locked descendants belong
    // to uid, then they are unlocked and X is locked instead.
    // Conditions to FAIL:
    //   1. X is already locked
    //   2. X has NO locked descendants
    //   3. Any locked descendant is NOT locked by uid
    //   4. Any ancestor of X is locked
    bool upgrade(const string& name, int uid) {
        TreeNode* node = nodeMap[name];

        // Check 1: X itself is locked?
        if (node->isLocked) return false;

        // Check 2: at least one locked descendant?
        if (node->lockedDescendants.empty()) return false;

        // Check 3: all locked descendants owned by same uid?
        for (TreeNode* desc : node->lockedDescendants) {
            if (desc->lockedBy != uid) return false;
        }

        // Check 4: any locked ancestor?
        // NOTE: This is technically implied by the fact that descendants
        // could only have been locked if ancestors were free. But we check
        // explicitly for safety.
        TreeNode* anc = node->parent;
        while (anc) {
            if (anc->isLocked) return false;
            anc = anc->parent;
        }

        // All checks passed — unlock all descendants, then lock X
        // Make a copy because we'll be modifying lockedDescendants during unlock
        set<TreeNode*> toUnlock = node->lockedDescendants;
        for (TreeNode* desc : toUnlock) {
            unlock(desc->name, uid);
        }
        lock(name, uid);
        return true;
    }
};

// ─── MAIN ─────────────────────────────────────────────────────────────────────
int main() {
    ios_base::sync_with_stdio(false);
    cin.tie(NULL);

    int N, M, Q;
    cin >> N >> M >> Q;

    vector<string> names(N);
    for (int i = 0; i < N; i++) cin >> names[i];

    MAryTree tree(names[0]);
    tree.build(names, M);

    while (Q--) {
        int type, uid;
        string nodeName;
        cin >> type >> nodeName >> uid;

        bool result = false;
        if      (type == 1) result = tree.lock(nodeName, uid);
        else if (type == 2) result = tree.unlock(nodeName, uid);
        else if (type == 3) result = tree.upgrade(nodeName, uid);

        cout << (result ? "true" : "false") << "\n";
    }

    return 0;
}

Complexity Analysis (Part A)

Operation	Time	Reason
`lock`	O(log_M N)	Walk ancestors to check; `lockedDescendants.size()` is O(1)
`unlock`	O(log_M N)	Walk ancestors to update their sets
`upgrade`	O(K · log_M N)	K = # locked descendants; unlock each in O(log_M N)
Build tree	O(N)	BFS traversal
Space	O(N)	One set per node

log_M N = height of balanced M-ary tree with N nodes

Part B Solution — Thread Safety Discussion

The Question (Actual Part B from Juspay Hackathon)

You are running lock() on a multi-core machine with multiple threads. Multiple threads may call lock() simultaneously on different nodes.

Will there be a race condition? If yes, how do you make it thread-safe?

Constraints:

Multiple threads running simultaneously should not affect correctness

Make critical sections as granular as possible (don't lock the whole tree)

Each individual memory operation (read/write of one variable) is atomic

Race Condition Analysis

Race Condition 1: Same node, two threads

Thread T1: lock("India", 5)
Thread T2: lock("India", 5)

Expected: one returns true, one returns false
Problem:  both may return true!

Timeline:
T1 checks: India.isLocked == false  ✅ passes
[CONTEXT SWITCH]
T2 checks: India.isLocked == false  ✅ passes (stale read!)
T1 sets:   India.isLocked = true    → returns true
T2 sets:   India.isLocked = true    → also returns true ❌ WRONG

Race Condition 2: Ancestor-Descendant conflict

Thread T1: lock("World", 1)
Thread T2: lock("Asia", 2)

Expected: one succeeds, one fails
Problem:  Both may succeed!

Timeline:
T1 checks: World.lockedDescendants.size() == 0 ✅
[CONTEXT SWITCH]
T2 enters while(par) loop, par = World
T2 checks: World.isLocked == false  ✅
[CONTEXT SWITCH]
T1 sets:   World.isLocked = true → returns true
T2 continues loop → par becomes null → sets Asia.isLocked = true → returns true ❌ WRONG

Race Condition 3: Concurrent ancestor set modification

Thread T1: lock("China", 1)
Thread T2: lock("India", 2)
Both walk up to Asia → Asia.lockedDescendants.insert() called simultaneously
→ set corruption (not thread-safe container)

Solution A: Coarse-grained Mutex (Naive — Not Recommended)

#include <mutex>
std::mutex globalMutex;

bool lock(const string& name, int uid) {
    std::lock_guard<std::mutex> guard(globalMutex);
    // ... same logic ...
}

Problem: Kills all parallelism. Only one thread can work at a time.

Solution B: Read-Write Lock Per Node (Granular Parallelism)

Key insight: Multiple threads can check ancestor lock states concurrently using read (shared) locks. Modifying a node's lock status or updating an ancestor's lockedDescendants set requires an exclusive (unique) write lock.

#include <shared_mutex>
#include <mutex>

struct TreeNode {
    string name;
    int lockedBy;
    bool isLocked;
    vector<TreeNode*> children;
    TreeNode* parent;
    set<TreeNode*> lockedDescendants;
    mutable shared_mutex lockMtx;   // Protects isLocked & lockedBy (R/W)
    mutex descendantsMtx;           // Protects lockedDescendants set mutations
};

bool lock(const string& name, int uid) {
    TreeNode* node = nodeMap[name];

    // 1. Exclusive WRITE lock on current node
    unique_lock<shared_mutex> nodeLock(node->lockMtx);

    {
        lock_guard<mutex> descGuard(node->descendantsMtx);
        if (node->isLocked || !node->lockedDescendants.empty()) return false;
    }

    // 2. READ lock on all ancestors (verifies none are locked)
    vector<shared_lock<shared_mutex>> ancestorLocks;
    TreeNode* anc = node->parent;
    while (anc) {
        ancestorLocks.emplace_back(anc->lockMtx);
        if (anc->isLocked) return false; // locks auto-released on scope exit
        anc = anc->parent;
    }

    // 3. Mark node as locked
    node->isLocked = true;
    node->lockedBy = uid;

    // 4. Update ancestors' lockedDescendants with mutex protection
    TreeNode* cur = node->parent;
    while (cur) {
        lock_guard<mutex> descGuard(cur->descendantsMtx);
        cur->lockedDescendants.insert(node);
        cur = cur->parent;
    }

    return true;
}

Why this achieves high concurrency for siblings:

T1 locking "China" and T2 locking "India" both acquire READ (shared) locks on their common ancestor Asia to verify Asia->isLocked == false. Both can read simultaneously.
When updating Asia->lockedDescendants, each thread briefly acquires Asia->descendantsMtx, safely inserting its node without corrupting the set.
If T3 tries to lock Asia itself, it attempts to acquire a WRITE (exclusive) lock on Asia->lockMtx, which correctly blocks until T1 and T2 finish their ancestor checks.
Deadlock prevention: Consistent lock acquisition order (always bottom-up from child to root).

Solution C: Lock-Free (Set-and-Test, No Mutex/Atomic)

The interviewer pushed for a solution without any mutex or atomic primitives. This is the approach hinted at in the hackathon:

Core Idea: Instead of isLocked (bool), use lockedCount (int).

"Test-and-Set" (buggy): Check if 0, then set to 1 — race between check and set.
"Set-and-Test" (safer): Increment first, then check if > 1 (someone else also incremented).

struct TreeNode {
    string name;
    int lockedBy;
    int lockedCount;  // 0 = free, 1 = locked by one thread
    vector<TreeNode*> children;
    TreeNode* parent;
    set<TreeNode*> lockedDescendants;
};

bool lock(const string& name, int uid) {
    TreeNode* node = nodeMap[name];

    // SET first, then TEST
    node->lockedCount++;                             // (A) increment
    if (node->lockedCount > 1 || !node->lockedDescendants.empty()) {
        node->lockedCount--;                         // rollback
        return false;
    }

    // Walk up to root, insert self into ancestor's lockedDescendants
    // and check that ancestor is not itself locked
    TreeNode* par = node->parent;
    while (par) {
        par->lockedDescendants.insert(node);         // SET before TEST
        if (par->lockedCount > 0) {                  // TEST after SET
            // Rollback: remove from all ancestors we already inserted into
            par->lockedDescendants.erase(node);
            node->lockedCount--;
            TreeNode* rollback = node->parent;
            while (rollback != par) {
                rollback->lockedDescendants.erase(node);
                rollback = rollback->parent;
            }
            return false;
        }
        par = par->parent;
    }

    node->lockedBy = uid;
    return true;
}

⚠️ As the original author notes in the interview transcript: Even this approach has race conditions because lockedCount++ itself is not atomic without hardware support. The interviewer's intended "solution" was acknowledged as theoretically flawed. The correct production answer is to use std::atomic<int> or a real mutex.

Solution D: Per-Node Atomic Counter (Correct Lock-Free)

#include <atomic>

struct TreeNode {
    string name;
    int lockedBy;
    std::atomic<int> lockedCount{0}; // atomic guarantees read-modify-write
    vector<TreeNode*> children;
    TreeNode* parent;
    // Use atomic operations for lockedDescendants too, or protect with node mutex
};

bool lock(const string& name, int uid) {
    TreeNode* node = nodeMap[name];

    int expected = 0;
    // compare_exchange: atomically set lockedCount from 0 to 1
    if (!node->lockedCount.compare_exchange_strong(expected, 1)) {
        return false; // someone else locked it first
    }

    // Now we have exclusive claim on this node
    // Check descendants (via lockedDescendants set — needs protection)
    // ...
}

Interview Follow-Up Q&A

Q1: Is the ancestor check in Upgrade redundant?

Almost — because a descendant can only be locked if all its ancestors were free at lock time. However, if between the time a descendant was locked and the upgrade attempt, an ancestor gets locked (via another branch), the check IS needed. So it's NOT fully redundant in a concurrent system.

Q2: Can we use unordered_set instead of set for lockedDescendants?

Yes. Time complexity stays the same asymptotically (O(1) average insert/erase vs O(log K)), but constant factors improve. The code runs faster in practice.

Q3: What's the critical section in lock()?

Two critical sections:

Reading and writing node->isLocked / lockedCount

Inserting into ancestor->lockedDescendants while walking up

Q4: Can two sibling nodes be locked concurrently without conflict?

Yes! Siblings share ancestors, but verifying ancestors only requires READ access on their lock status (isLocked). Using shared_mutex for lock status and granular mutexes for ancestor descendant sets, multiple threads can check ancestors and update descendant sets concurrently without blocking each other.

🏁 Final Checklist Before Submitting

Tree builds correctly in BFS order
unordered_map<string, TreeNode*> for O(1) node lookup
lockedDescendants set maintained on every lock/unlock
propagateUp() walks all the way to root (not just parent)
Upgrade copies lockedDescendants before iterating (modifying while iterating = UB)
Unlock rolls back parent sets correctly
Edge cases: leaf node upgrade, lock already-locked node, wrong uid unlock

tree_of_space_solution.md