Problem Statement
Given a binary matrix mat[][] containing only 0s and 1s, and an integer k, we are also given multiple queries.
Each query contains a cell (i, j) that represents the center of a square.
For every query, we need to find the largest odd-sized square centered at (i, j) such that the square contains at most k ones.
If even the smallest 1 × 1 square contains more than k ones, the answer is -1.
Understanding the Square
Because the square must be centered at (i, j) and expand equally in all four directions, its side length is always odd.
For example:
Radius = 0
1 × 1
Radius = 1
3 × 3
Radius = 2
5 × 5
In general:
Side Length = 2 × radius + 1
So instead of directly searching for the side length, we can search for the maximum possible radius.
Example
Consider:
1 0 1 0 0
1 0 1 1 1
1 1 1 1 1
1 0 0 1 0
Suppose the query is:
(1, 2)
and:
k = 9
The center is:
row = 1
column = 2
The possible squares are:
Radius 0
1 × 1
Only the center cell is included.
Radius 1
3 × 3
Rows 0 to 2 and columns 1 to 3 are included.
This square contains 6 ones.
Since:
6 <= 9
the square is valid.
Radius 2
A 5 × 5 square would be required, but the center is too close to the matrix boundary.
Therefore, the largest valid square is:
3 × 3
So the answer is:
3
Main Idea
There are two important techniques used to solve this problem efficiently:
2D Prefix Sum
Binary Search
The constraints allow up to:
n, m <= 500
queries <= 10^4
Checking every cell for every possible square would be too expensive.
Instead, we first preprocess the matrix using a 2D prefix sum.
This allows us to find the number of ones inside any rectangular region in O(1) time.
Then, for each query, we use binary search to find the largest valid square.
1. Building the 2D Prefix Sum
We create:
int[][] prefix = new int[n + 1][m + 1];
The extra row and column make the calculations easier.
The formula is:
prefix[i][j] =
mat[i - 1][j - 1]
+ prefix[i - 1][j]
+ prefix[i][j - 1]
- prefix[i - 1][j - 1];
Why do we subtract:
prefix[i - 1][j - 1]
Because that area was counted twice.
Visually:
prefix[i-1][j]
|
↓
+-------+-------+
| | |
| A | B |
| | |
+-------+-------+
| | |
| C | D |
| | |
+-------+-------+
↑
prefix[i][j-1]
When we add the top and left portions, the top-left portion is included twice, so we subtract it once.
2. Getting the Number of Ones in a Square
Suppose the square boundaries are:
top
bottom
left
right
We can calculate the number of ones using:
private int getSum(
int[][] prefix,
int top,
int left,
int bottom,
int right) {
return prefix[bottom + 1][right + 1]
- prefix[top][right + 1]
- prefix[bottom + 1][left]
+ prefix[top][left];
}
This gives the number of ones inside:
[top ... bottom]
[left ... right]
in constant time:
O(1)
3. Finding the Maximum Possible Radius
For a query:
int r = query[0];
int c = query[1];
the square must remain inside the matrix.
Therefore, the radius is limited by the distance from the center to each boundary.
The four distances are:
distance to top = r
distance to bottom = n - 1 - r
distance to left = c
distance to right = m - 1 - c
The smallest of these determines the maximum radius.
So:
int maxRadius = Math.min(
Math.min(r, n - 1 - r),
Math.min(c, m - 1 - c)
);
For example, if:
r = 2
c = 3
and the matrix is 5 × 7, then:
top = 2
bottom = 2
left = 3
right = 3
The smallest distance is 2, so the maximum radius is 2.
4. Important Edge Case
Before performing binary search, we check the 1 × 1 square.
if (mat[r][c] > k) {
ans.add(-1);
continue;
}
Why?
The smallest possible square contains only the center cell.
If the center contains 1 and:
k = 0
then:
1 > 0
Therefore, even the smallest square is invalid.
There is no answer, so we return:
-1
This is an important case that can otherwise cause a wrong answer.
5. Binary Search
Now we know:
We binary search between:
int low = 0;
int high = maxRadius;
We calculate:
int mid = low + (high - low) / 2;
For this radius, the square boundaries are:
int top = r - mid;
int bottom = r + mid;
int left = c - mid;
int right = c + mid;
Then we calculate the number of ones:
int ones = getSum(
prefix,
top,
left,
bottom,
right
);
Why Can We Use Binary Search?
This is the most important observation.
Suppose a square with radius 2 contains 7 ones.
If we increase the radius to 3, the new square contains everything from the radius 2 square plus some additional cells.
Therefore, the number of ones can never decrease.
For example:
Radius 0 → 2 ones
Radius 1 → 4 ones
Radius 2 → 7 ones
Radius 3 → 12 ones
If:
k = 8
then:
radius 0 → valid
radius 1 → valid
radius 2 → valid
radius 3 → invalid
Once a radius becomes invalid, every larger radius will also be invalid.
This gives us a monotonic property:
Valid → Valid → Valid → Invalid → Invalid
Therefore, binary search can find the largest valid radius.
6. Binary Search Logic
If the number of ones is at most k:
if (ones <= k) {
best = mid;
low = mid + 1;
}
The current radius is valid, so we save it.
Then we try to find a larger square:
low = mid + 1;
Otherwise:
else {
high = mid - 1;
}
The current square contains too many ones, so we need a smaller radius.
7. Converting Radius to Side Length
Once binary search finishes, best contains the largest valid radius.
The side length is:
2 × radius + 1
Therefore:
ans.add(2 * best + 1);
For example:
radius = 0 → 1
radius = 1 → 3
radius = 2 → 5
radius = 3 → 7
Complete Java Solution
import java.util.*;
class Solution {
ArrayList<Integer> largestSquare(int[][] mat, int[][] queries, int k) {
int n = mat.length;
int m = mat[0].length;
// Build 2D prefix sum
int[][] prefix = new int[n + 1][m + 1];
for (int i = 1; i <= n; i++) {
for (int j = 1; j <= m; j++) {
prefix[i][j] =
mat[i - 1][j - 1]
+ prefix[i - 1][j]
+ prefix[i][j - 1]
- prefix[i - 1][j - 1];
}
}
ArrayList<Integer> ans = new ArrayList<>();
for (int[] query : queries) {
int r = query[0];
int c = query[1];
// Maximum possible radius based on boundaries
int maxRadius = Math.min(
Math.min(r, n - 1 - r),
Math.min(c, m - 1 - c)
);
// Check 1 x 1 square
if (mat[r][c] > k) {
ans.add(-1);
continue;
}
int low = 0;
int high = maxRadius;
int best = 0;
while (low <= high) {
int mid = low + (high - low) / 2;
int top = r - mid;
int bottom = r + mid;
int left = c - mid;
int right = c + mid;
int ones = getSum(
prefix,
top,
left,
bottom,
right
);
if (ones <= k) {
// Current square is valid
best = mid;
// Try a larger square
low = mid + 1;
} else {
// Too many ones
high = mid - 1;
}
}
// Convert radius to side length
ans.add(2 * best + 1);
}
return ans;
}
private int getSum(
int[][] prefix,
int top,
int left,
int bottom,
int right) {
return prefix[bottom + 1][right + 1]
- prefix[top][right + 1]
- prefix[bottom + 1][left]
+ prefix[top][left];
}
}
Dry Run
Consider:
mat =
1 0 1 0 0
1 0 1 1 1
1 1 1 1 1
1 0 0 1 0
Query:
[1, 2]
and:
k = 9
Maximum radius:
r = 1
c = 2
Distances:
top = 1
bottom = 2
left = 2
right = 2
Therefore:
maxRadius = 1
Binary search checks:
radius = 0
Square:
1 × 1
Number of ones:
1
Valid.
Then:
radius = 1
Square:
3 × 3
Number of ones:
6
Valid because:
6 <= 9
Therefore:
best = 1
Convert radius to side length:
2 × 1 + 1 = 3
Answer:
3
Complexity Analysis
Let:
n = number of rows
m = number of columns
q = number of queries
Prefix Sum
We visit every matrix cell once:
O(n × m)
Each Query
Binary search checks at most:
O(log(min(n, m)))
radii.
Each radius check takes:
O(1)
because of the prefix sum.
Therefore:
O(q × log(min(n, m)))
for all queries.
Total Time Complexity
O(n × m + q × log(min(n, m)))
Space Complexity
The prefix sum matrix requires:
O(n × m)
space.
Key Interview Takeaways
This problem combines three important concepts:
1. Prefix Sum
Use a 2D prefix sum when you need to repeatedly calculate the sum of arbitrary rectangular regions.
2. Binary Search on Answer
The answer is not directly searchable, but the possible radii have a monotonic property:
valid → valid → valid → invalid → invalid
That makes binary search possible.
3. Center-Based Expansion
For an odd square centered at (r, c):
top = r - radius
bottom = r + radius
left = c - radius
right = c + radius
and:
side = 2 × radius + 1
The combination of these three ideas reduces a potentially expensive matrix-search problem to the required:
O(n × m + q × log(min(n, m)))
which is efficient for n, m <= 500 and up to 10^4 queries.