Design a task executor that uses fine-grained locking to coordinate operations by key.

Implement the KeyedTaskExecutor class:

KeyedTaskExecutor(keyCount) creates an executor for keys from 0 through keyCount - 1.
execute(key, task) waits for exclusive access to key, then invokes task exactly once before returning.

Tasks submitted with the same key must never overlap. Tasks using different keys must be independent and capable of executing at the same time; an operation on one key must not be blocked solely because another key is busy.

The key's lock must remain held for the callback's entire execution. If task throws an unchecked exception or panics, the lock must still be released and the original failure must continue to the caller.

Different KeyedTaskExecutor instances must also be independent. No fairness or same-key ordering policy is required.

The judge creates all caller threads and supplies the callbacks. Standard concurrency, callback, lock, and thread APIs are preloaded, so you do not need import, include, package, or using statements.

Example 1:

Input:

Output:

Explanation: Both operations use key 1, so they share the same lock. Either task may run first.

Example 2:

Input:

Output:

Explanation: Keys 0 and 2 have separate locks, so neither operation has to wait for the other.

Constraints

1 <= keyCount <= 100
0 <= key < keyCount
At most 100 threads use one executor at a time.
At most 50000 calls are made per executor.
A callback may complete normally or fail with an unchecked exception or panic.
A callback does not call execute on the same executor.
Judge threads are not interrupted while waiting for a key.

Starter Code

Understanding the Problem

A coarse-grained design could protect every task with one global lock. That would correctly serialize same-key operations, but it would also make unrelated keys wait for one another.

A fine-grained design associates a separate lock with each key. Operations contend only when they access the same logical resource. This preserves mutual exclusion where it is needed while allowing unrelated work to proceed concurrently.

The callback defines the critical section. Releasing the key's lock before invoking task would leave the actual operation unprotected.

Synchronization Strategy

Create an array or list containing one lock per key.
In execute(key, task), select the lock at index key.
Acquire that lock.
Invoke the callback while the lock remains held.
Release the lock in an exception-safe cleanup path.

Two calls with the same key select the same lock and therefore cannot overlap. Calls with different keys select different locks, so they do not block one another.

The lock collection belongs to the executor instance. Making it static or global would incorrectly couple separate executors.

Exception-Safe Release

Every successful lock acquisition must have one matching release, even when the callback fails.

Java uses try/finally, Python uses a with statement, C++ uses lock_guard, Go uses defer, and C# uses lock. These mechanisms release the selected key's lock while preserving the callback's original exception or panic.

Correctness

For a fixed key, every call acquires the same exclusive lock before invoking its callback. At most one callback for that key can therefore execute at a time.

For two different keys, the calls acquire different locks. Holding one key's lock does not prevent acquisition of the other, so the callbacks can execute concurrently.

The lock remains held until the callback returns or fails. Thus, the entire operation is protected, and exception-safe cleanup allows later work on that key to continue.

Because every executor constructs its own lock collection, operations on different instances are independent even when they use the same numeric key.

Code

Complexity Analysis

Constructor: O(keyCount) time and space.
execute: O(1) synchronization work plus the callback's running time.

A call may block behind earlier work using the same key, but work on other keys does not contribute to that wait.

Design a Keyed Task Executor