3391. Design a 3D Binary Matrix with Efficient Layer Tracking

MediumDesignArrayHash TableMatrixOrdered SetHeap (Priority Queue)

Problem Description

You need to implement a class that manages a 3D binary matrix (a cube) of size n x n x n where each cell can contain either 0 or 1.

The class Matrix3D should support the following operations:

Constructor Matrix3D(int n): Creates a 3D binary matrix of dimensions n x n x n with all cells initially set to 0.
setCell(int x, int y, int z): Sets the cell at position (x, y, z) to 1. The position is specified by:
- x: the layer index (0 to n-1)
- y: the row index within layer x (0 to n-1)
- z: the column index within row y (0 to n-1)
unsetCell(int x, int y, int z): Sets the cell at position (x, y, z) to 0.
largestMatrix(): Returns the layer index x that contains the maximum number of 1s. If multiple layers have the same maximum count of 1s, return the largest layer index among them.

For example, in a 3x3x3 matrix, you have 3 layers (x = 0, 1, 2), and each layer is a 3x3 2D matrix. The largestMatrix() function counts how many 1s are in each layer and returns the layer with the most 1s. If layer 0 has 5 ones and layer 2 has 5 ones, it would return 2 (the larger index).

Quick Interview Experience

Help others by sharing your interview experience

Have you seen this problem before?

Intuition

The key insight is that we need to efficiently track which layer (2D plane at index x) has the most 1s, and handle ties by returning the largest index.

Since we're frequently updating individual cells and querying for the layer with maximum 1s, we need to avoid recounting all cells every time largestMatrix() is called. Instead, we can maintain a running count of 1s for each layer.

We maintain:

A 3D array g to store the actual matrix values
An array cnt where cnt[x] tracks how many 1s are in layer x
When setting a cell from 0 to 1, we increment cnt[x] by 1
When unsetting a cell from 1 to 0, we decrement cnt[x] by 1

However, finding the maximum in cnt array each time largestMatrix() is called would still be O(n). To optimize this further, we need a data structure that maintains the layers sorted by their count of 1s.

This leads us to use a SortedList (ordered set) that stores tuples (count, layer_index). By using a custom sorting key (-count, -layer_index), we ensure:

Layers with more 1s appear first (negative count for descending order)
Among layers with equal counts, larger indices appear first (negative index for descending order)

With this structure, largestMatrix() becomes O(1) - just return the layer index from the first element in the sorted list. The trade-off is that setCell and unsetCell operations now have O(log n) complexity for updating the sorted list, but this is acceptable since we optimize the frequent query operation.

The solution also handles edge cases by checking if a cell is already set/unset before modifying it, preventing incorrect count updates.

Learn more about Heap (Priority Queue) patterns.

Solution Implementation

1from sortedcontainers import SortedList 2 3 4class matrix3D: 5 def __init__(self, n: int): 6 """ 7 Initialize a 3D matrix of size n x n x n. 8 9 Args: 10 n: The dimension of the cubic matrix 11 """ 12 # Create a 3D matrix filled with zeros 13 self.grid = [[[0] * n for _ in range(n)] for _ in range(n)] 14 15 # Track the count of set cells for each x-plane (layer) 16 self.cell_count = [0] * n 17 18 # Maintain a sorted list of (count, x_index) tuples 19 # Sorted by count (descending) then by x_index (descending) 20 self.sorted_planes = SortedList(key=lambda item: (-item[0], -item[1])) 21 22 def setCell(self, x: int, y: int, z: int) -> None: 23 """ 24 Set the cell at position (x, y, z) to 1. 25 26 Args: 27 x: The x-coordinate (plane index) 28 y: The y-coordinate 29 z: The z-coordinate 30 """ 31 # Check if cell is already set 32 if self.grid[x][y][z] == 1: 33 return 34 35 # Set the cell to 1 36 self.grid[x][y][z] = 1 37 38 # Remove the old count entry for this x-plane from sorted list 39 self.sorted_planes.discard((self.cell_count[x], x)) 40 41 # Increment the count for this x-plane 42 self.cell_count[x] += 1 43 44 # Add the updated count entry back to sorted list 45 self.sorted_planes.add((self.cell_count[x], x)) 46 47 def unsetCell(self, x: int, y: int, z: int) -> None: 48 """ 49 Set the cell at position (x, y, z) to 0. 50 51 Args: 52 x: The x-coordinate (plane index) 53 y: The y-coordinate 54 z: The z-coordinate 55 """ 56 # Check if cell is already unset 57 if self.grid[x][y][z] == 0: 58 return 59 60 # Set the cell to 0 61 self.grid[x][y][z] = 0 62 63 # Remove the old count entry for this x-plane from sorted list 64 self.sorted_planes.discard((self.cell_count[x], x)) 65 66 # Decrement the count for this x-plane 67 self.cell_count[x] -= 1 68 69 # Only add back to sorted list if count is non-zero 70 if self.cell_count[x] > 0: 71 self.sorted_planes.add((self.cell_count[x], x)) 72 73 def largestMatrix(self) -> int: 74 """ 75 Return the index of the x-plane with the most set cells. 76 If no cells are set, return n-1 (the last plane index). 77 78 Returns: 79 The x-index of the plane with the most set cells 80 """ 81 # Return the x-index from the first entry (highest count) 82 # If sorted list is empty, return the last valid index 83 if self.sorted_planes: 84 return self.sorted_planes[0][1] 85 else: 86 return len(self.grid) - 1 87 88 89# Your matrix3D object will be instantiated and called as such: 90# obj = matrix3D(n) 91# obj.setCell(x,y,z) 92# obj.unsetCell(x,y,z) 93# param_3 = obj.largestMatrix() 94

1/** 2 * A 3D matrix data structure that tracks set cells and maintains 3 * statistics about which x-plane has the most set cells. 4 */ 5class matrix3D { 6 // 3D array to store the state of each cell (0 or 1) 7 private final int[][][] grid; 8 9 // Array to count the number of set cells in each x-plane 10 private final int[] xPlaneCounts; 11 12 // TreeSet to maintain x-planes sorted by their cell count (descending) 13 // and by x-index (descending) as tiebreaker 14 private final TreeSet<int[]> sortedPlanes; 15 16 /** 17 * Initializes a 3D matrix of size n x n x n. 18 * @param n The dimension of the cubic matrix 19 */ 20 public matrix3D(int n) { 21 grid = new int[n][n][n]; 22 xPlaneCounts = new int[n]; 23 24 // Custom comparator: sort by count (descending), then by index (descending) 25 sortedPlanes = new TreeSet<>((a, b) -> { 26 if (a[0] == b[0]) { 27 return b[1] - a[1]; // If counts are equal, sort by index descending 28 } 29 return b[0] - a[0]; // Sort by count descending 30 }); 31 } 32 33 /** 34 * Sets the cell at position (x, y, z) to 1. 35 * Updates the count for x-plane and maintains sorted order. 36 * @param x The x-coordinate 37 * @param y The y-coordinate 38 * @param z The z-coordinate 39 */ 40 public void setCell(int x, int y, int z) { 41 // Skip if cell is already set 42 if (grid[x][y][z] == 1) { 43 return; 44 } 45 46 // Set the cell 47 grid[x][y][z] = 1; 48 49 // Update the sorted set: remove old count entry 50 sortedPlanes.remove(new int[] {xPlaneCounts[x], x}); 51 52 // Increment count for this x-plane 53 xPlaneCounts[x]++; 54 55 // Add updated count entry back to sorted set 56 sortedPlanes.add(new int[] {xPlaneCounts[x], x}); 57 } 58 59 /** 60 * Unsets the cell at position (x, y, z) to 0. 61 * Updates the count for x-plane and maintains sorted order. 62 * @param x The x-coordinate 63 * @param y The y-coordinate 64 * @param z The z-coordinate 65 */ 66 public void unsetCell(int x, int y, int z) { 67 // Skip if cell is already unset 68 if (grid[x][y][z] == 0) { 69 return; 70 } 71 72 // Unset the cell 73 grid[x][y][z] = 0; 74 75 // Update the sorted set: remove old count entry 76 sortedPlanes.remove(new int[] {xPlaneCounts[x], x}); 77 78 // Decrement count for this x-plane 79 xPlaneCounts[x]--; 80 81 // Only add back to sorted set if plane still has set cells 82 if (xPlaneCounts[x] > 0) { 83 sortedPlanes.add(new int[] {xPlaneCounts[x], x}); 84 } 85 } 86 87 /** 88 * Returns the x-index of the plane with the most set cells. 89 * If multiple planes have the same count, returns the largest x-index. 90 * If no cells are set, returns n-1 (the last plane index). 91 * @return The x-index of the plane with the most set cells 92 */ 93 public int largestMatrix() { 94 if (sortedPlanes.isEmpty()) { 95 return grid.length - 1; // Return last plane index if no cells are set 96 } 97 return sortedPlanes.first()[1]; // Return x-index of plane with most cells 98 } 99} 100 101/** 102 * Your matrix3D object will be instantiated and called as such: 103 * matrix3D obj = new matrix3D(n); 104 * obj.setCell(x,y,z); 105 * obj.unsetCell(x,y,z); 106 * int param_3 = obj.largestMatrix(); 107 */ 108

1class matrix3D { 2private: 3 // 3D matrix to store cell states (0 or 1) 4 vector<vector<vector<int>>> grid; 5 6 // Count of set cells in each x-layer (plane at x=i) 7 vector<int> layerCounts; 8 9 // Set to maintain x-layers sorted by count (descending) and x-index 10 // Stored as {-count, -x} to achieve descending order of counts 11 set<pair<int, int>> sortedLayers; 12 13public: 14 matrix3D(int n) { 15 // Initialize n×n×n grid with all cells unset (0) 16 grid.resize(n, vector<vector<int>>(n, vector<int>(n, 0))); 17 18 // Initialize count for each x-layer to 0 19 layerCounts.resize(n, 0); 20 } 21 22 void setCell(int x, int y, int z) { 23 // Skip if cell is already set 24 if (grid[x][y][z] == 1) { 25 return; 26 } 27 28 // Mark the cell as set 29 grid[x][y][z] = 1; 30 31 // Update the sorted set: remove old entry for layer x 32 sortedLayers.erase({-layerCounts[x], -x}); 33 34 // Increment count for layer x 35 layerCounts[x]++; 36 37 // Insert updated entry for layer x 38 sortedLayers.insert({-layerCounts[x], -x}); 39 } 40 41 void unsetCell(int x, int y, int z) { 42 // Skip if cell is already unset 43 if (grid[x][y][z] == 0) { 44 return; 45 } 46 47 // Mark the cell as unset 48 grid[x][y][z] = 0; 49 50 // Update the sorted set: remove old entry for layer x 51 sortedLayers.erase({-layerCounts[x], -x}); 52 53 // Decrement count for layer x 54 layerCounts[x]--; 55 56 // Only insert back if layer still has set cells 57 if (layerCounts[x] > 0) { 58 sortedLayers.insert({-layerCounts[x], -x}); 59 } 60 } 61 62 int largestMatrix() { 63 // If no layers have set cells, return the last layer index 64 // Otherwise, return the x-index of the layer with most set cells 65 return sortedLayers.empty() ? grid.size() - 1 : -sortedLayers.begin()->second; 66 } 67}; 68 69/** 70 * Your matrix3D object will be instantiated and called as such: 71 * matrix3D* obj = new matrix3D(n); 72 * obj->setCell(x,y,z); 73 * obj->unsetCell(x,y,z); 74 * int param_3 = obj->largestMatrix(); 75 */ 76

1// 3D matrix to store cell states (0 or 1) 2let grid: number[][][]; 3 4// Count of set cells in each x-layer (plane at x=i) 5let layerCounts: number[]; 6 7// Map to maintain x-layers sorted by count (descending) and x-index 8// Using a Map with composite key as string "count,x" for sorting 9let sortedLayers: Map<string, [number, number]>; 10 11function matrix3D(n: number): void { 12 // Initialize n×n×n grid with all cells unset (0) 13 grid = Array(n).fill(null).map(() => 14 Array(n).fill(null).map(() => 15 Array(n).fill(0) 16 ) 17 ); 18 19 // Initialize count for each x-layer to 0 20 layerCounts = Array(n).fill(0); 21 22 // Initialize the sorted layers map 23 sortedLayers = new Map(); 24} 25 26function setCell(x: number, y: number, z: number): void { 27 // Skip if cell is already set 28 if (grid[x][y][z] === 1) { 29 return; 30 } 31 32 // Mark the cell as set 33 grid[x][y][z] = 1; 34 35 // Remove old entry for layer x from sorted map if it exists 36 const oldKey = `${-layerCounts[x]},${-x}`; 37 sortedLayers.delete(oldKey); 38 39 // Increment count for layer x 40 layerCounts[x]++; 41 42 // Insert updated entry for layer x with negative values for descending sort 43 const newKey = `${-layerCounts[x]},${-x}`; 44 sortedLayers.set(newKey, [-layerCounts[x], -x]); 45 46 // Sort the map entries to maintain order 47 sortedLayers = new Map([...sortedLayers.entries()].sort()); 48} 49 50function unsetCell(x: number, y: number, z: number): void { 51 // Skip if cell is already unset 52 if (grid[x][y][z] === 0) { 53 return; 54 } 55 56 // Mark the cell as unset 57 grid[x][y][z] = 0; 58 59 // Remove old entry for layer x from sorted map 60 const oldKey = `${-layerCounts[x]},${-x}`; 61 sortedLayers.delete(oldKey); 62 63 // Decrement count for layer x 64 layerCounts[x]--; 65 66 // Only insert back if layer still has set cells 67 if (layerCounts[x] > 0) { 68 const newKey = `${-layerCounts[x]},${-x}`; 69 sortedLayers.set(newKey, [-layerCounts[x], -x]); 70 71 // Sort the map entries to maintain order 72 sortedLayers = new Map([...sortedLayers.entries()].sort()); 73 } 74} 75 76function largestMatrix(): number { 77 // If no layers have set cells, return the last layer index 78 // Otherwise, return the x-index of the layer with most set cells 79 if (sortedLayers.size === 0) { 80 return grid.length - 1; 81 } 82 83 // Get the first entry (layer with most set cells due to sorting) 84 const firstEntry = sortedLayers.entries().next().value; 85 return -firstEntry[1][1]; 86} 87

Solution Approach

The implementation uses a combination of a 3D array for storage and a sorted list for efficient querying:

Data Structures:

3D Array g: A n x n x n array where g[x][y][z] stores the binary value (0 or 1) at position (x, y, z).
Count Array cnt: An array of length n where cnt[x] tracks the total number of 1s in layer x.
Sorted List sl: An ordered set that maintains tuples (cnt[x], x) sorted by:
- First by count in descending order (more 1s come first)
- Then by layer index in descending order (larger indices come first for ties)
This is achieved using the key function: lambda x: (-x[0], -x[1])

Implementation Details:

__init__(n):

Initialize the 3D array g with all zeros: [[[0] * n for _ in range(n)] for _ in range(n)]
Initialize count array cnt with all zeros: [0] * n
Create an empty SortedList with the custom sorting key

setCell(x, y, z):

Check if g[x][y][z] is already 1; if so, return immediately (avoid duplicate counting)
Set g[x][y][z] = 1
Remove the old tuple (cnt[x], x) from the sorted list using discard()
Increment cnt[x] by 1
Add the updated tuple (cnt[x], x) back to the sorted list

unsetCell(x, y, z):

Check if g[x][y][z] is already 0; if so, return immediately
Set g[x][y][z] = 0
Remove the old tuple (cnt[x], x) from the sorted list
Decrement cnt[x] by 1
Only add (cnt[x], x) back to the sorted list if cnt[x] > 0 (layers with no 1s are excluded)

largestMatrix():

If the sorted list is not empty, return sl[0][1] (the layer index from the first element)
If the sorted list is empty (all layers have 0 ones), return n - 1 (the largest possible layer index)

Time Complexity:

setCell and unsetCell: O(log n) due to sorted list operations
largestMatrix: O(1) since we just access the first element
Space Complexity: O(n³) for the 3D array plus O(n) for auxiliary structures

Ready to land your dream job?

Unlock your dream job with a 5-minute evaluator for a personalized learning plan!

Start Evaluator

Example Walkthrough

Let's walk through a small example with a 3x3x3 matrix to illustrate how the solution works.

Initial State:

Create Matrix3D(3) - a 3x3x3 cube with all zeros
g = 3D array of all zeros
cnt = [0, 0, 0] (no 1s in any layer)
sl = [] (empty sorted list)

Step 1: setCell(0, 1, 1)

Check: g[0][1][1] == 0? Yes, proceed
Set g[0][1][1] = 1
Remove (0, 0) from sorted list (using discard, no error if not present)
Update cnt[0] = 0 + 1 = 1
Add (1, 0) to sorted list
State: cnt = [1, 0, 0], sl = [(1, 0)]

Step 2: setCell(2, 0, 0)

Check: g[2][0][0] == 0? Yes, proceed
Set g[2][0][0] = 1
Remove (0, 2) from sorted list (not present, no effect)
Update cnt[2] = 0 + 1 = 1
Add (1, 2) to sorted list
State: cnt = [1, 0, 1], sl = [(1, 2), (1, 0)] (sorted by (-count, -index))

Step 3: largestMatrix()

Check first element: sl[0] = (1, 2)
Return layer index: 2 (layer 2 chosen over layer 0 due to larger index)

Step 4: setCell(0, 2, 2)

Check: g[0][2][2] == 0? Yes, proceed
Set g[0][2][2] = 1
Remove (1, 0) from sorted list
Update cnt[0] = 1 + 1 = 2
Add (2, 0) to sorted list
State: cnt = [2, 0, 1], sl = [(2, 0), (1, 2)] (layer 0 now has most 1s)

Step 5: largestMatrix()

Check first element: sl[0] = (2, 0)
Return layer index: 0 (layer 0 has 2 ones, more than layer 2's 1 one)

Step 6: unsetCell(0, 1, 1)

Check: g[0][1][1] == 1? Yes, proceed
Set g[0][1][1] = 0
Remove (2, 0) from sorted list
Update cnt[0] = 2 - 1 = 1
Add (1, 0) back to sorted list (since cnt[0] > 0)
State: cnt = [1, 0, 1], sl = [(1, 2), (1, 0)] (back to tie situation)

Step 7: largestMatrix()

Check first element: sl[0] = (1, 2)
Return layer index: 2 (tie broken by larger index)

This example demonstrates:

How the sorted list maintains layers ordered by count (descending) and index (descending for ties)
How setCell/unsetCell operations update both the matrix and the sorted list
How largestMatrix() returns the correct layer in O(1) time by checking the first element
How ties are handled by returning the largest layer index

Time and Space Complexity

Time Complexity:

__init__: O(n^3) for initializing the 3D matrix, plus O(1) for initializing the count array and SortedList, resulting in O(n^3) overall.
setCell: O(log n) - The method checks if a cell is already set in O(1), then performs discard and add operations on the SortedList, each taking O(log n) time.
unsetCell: O(log n) - Similar to setCell, it checks the cell value in O(1), then performs discard operation in O(log n), and conditionally adds to the SortedList in O(log n).
largestMatrix: O(1) - Simply accesses the first element of the SortedList and returns a value.

Space Complexity:

The 3D matrix self.g requires O(n^3) space to store n × n × n elements.
The count array self.cnt requires O(n) space for storing counts of each x-plane.
The SortedList self.sl stores at most n tuples, requiring O(n) space.
Overall space complexity: O(n^3) dominated by the 3D matrix storage.

Learn more about how to find time and space complexity quickly.

Common Pitfalls

1. Forgetting to Check if Cell State Has Already Changed

Pitfall: Not checking whether a cell is already set before setting it (or already unset before unsetting it) leads to incorrect counting. For example:

Calling setCell(0, 0, 0) twice would incorrectly increment cnt[0] twice
This breaks the invariant that cnt[x] should equal the actual number of 1s in layer x

Solution: Always check the current state before modifying:

def setCell(self, x: int, y: int, z: int) -> None:  if self.grid[x][y][z] == 1: # Already set, do nothing  return  # ... proceed with setting

2. Adding Zero-Count Entries to Sorted List

Pitfall: After unsetting cells, if a layer's count drops to 0, keeping (0, x) in the sorted list causes issues:

When all layers have 0 ones, the sorted list would contain multiple (0, x) entries
largestMatrix() would return an arbitrary layer index instead of the expected n-1

Solution: Only add non-zero counts to the sorted list:

def unsetCell(self, x: int, y: int, z: int) -> None:  # ... after decrementing count  if self.cell_count[x] > 0: # Only add if count is positive  self.sorted_planes.add((self.cell_count[x], x))

3. Incorrect Sorting Order for Tie-Breaking

Pitfall: Using the wrong sorting key leads to incorrect results when multiple layers have the same count. The problem requires returning the largest index among tied layers, but default sorting would return the smallest.

Solution: Use negative values in the sort key to reverse both orderings:

self.sorted_planes = SortedList(key=lambda item: (-item[0], -item[1])) # First element: -count (larger counts first) # Second element: -index (larger indices first for ties)

4. Not Handling Empty Sorted List in largestMatrix()

Pitfall: When all cells are 0 (sorted list is empty), directly accessing sl[0] throws an IndexError.

Solution: Check if the sorted list is non-empty before accessing:

def largestMatrix(self) -> int:  if self.sorted_planes:  return self.sorted_planes[0][1]  else:  return len(self.grid) - 1 # Return largest valid index

5. Using `remove()` Instead of `discard()` for SortedList Updates

Pitfall: Using remove() when the element might not exist (e.g., first time setting a cell in a layer) raises a ValueError.

Solution: Use discard() which silently does nothing if the element doesn't exist:

# Safe - won't raise error if (cnt[x], x) doesn't exist self.sorted_planes.discard((self.cell_count[x], x))

Discover Your Strengths and Weaknesses: Take Our 5-Minute Quiz to Tailor Your Study Plan:

In a binary min heap, the minimum element can be found in:

O(1)

O(logN)

O(N)

O(NlogN)

I don't know