Sprinkler square

The pump has one daily output and the orchard is a grid. The only question is how wide a square it can keep alive.

The problem

An orchard is planted in a rectangular grid. The cell grove[r][c] holds the litres the tree at row r, column c drinks each day; a bare plot is 0.

One sprinkler zone covers a square block — k rows by k columns, aligned to the grid, anywhere inside it. The pump delivers budget litres a day, so a zone is viable only if its trees need at most that.

Report the largest k with a viable zone, or 0 when even the cheapest tree is beyond the pump.

Input. grove — a list of rows, each a list of non-negative integers. budget — the pump's daily litres.

Output. The side length of the largest square block whose total demand is at most budget, or 0.

Example.

grove = [[1, 1, 3, 2],
         [4, 2, 1, 1],
         [3, 2, 2, 1],
         [1, 1, 1, 1]], budget = 8   ->  2

Rows 0–1, columns 0–1 need 1 + 1 + 4 + 2 = 8, exactly the pump's output. Every 3 by 3 block needs at least 12, so 3 is out of reach.

A second example, where the pump barely covers one tree:

grove = [[1, 1, 3, 2],
         [4, 2, 1, 1],
         [3, 2, 2, 1],
         [1, 1, 1, 1]], budget = 3   ->  1

The cheapest 2 by 2 block anywhere costs 5 litres, so only one tree fits.

Constraints.

• 1 <= len(grove) <= 300 and 1 <= len(grove[0]) <= 300
• 0 <= grove[r][c] <= 10^4
• 0 <= budget <= 10^9
• every row has the same length

Hints

Approach

Brute force

Sum every square at every corner and side. A 300 by 300 grove holds 2.7 · 10⁷ squares, a large one costing 9 · 10⁴ additions: near 10¹² operations.

The insight

Viability is monotone in the side length: shrinking a viable block by one row and one column can only lower its demand, so if side k works every smaller side works too.

Every tree needs zero or more litres, so a sub-block never totals more than the block enclosing it. The answers to "is side k viable?" read T T T F F as k grows, and the target is the last true. A prefix-sum table, pre[r][c] totalling the rectangle above and left of (r, c), reads any block in three additions, so one sweep of the corners settles a side in O(R · C).

Algorithm

1. Build pre, an (R+1) by (C+1) table with a zero row and column and pre[r+1][c+1] = grove[r][c] + pre[r][c+1] + pre[r+1][c] - pre[r][c].
2. Write fits(k): read every block of side k by inclusion–exclusion, true at the first total within budget.
3. Binary search k over [0, min(R, C)] for the largest true, with mid = (lo + hi + 1) // 2 and lo = mid when fits(mid).
4. Return lo.

Complexity

Time O(R · C · log min(R, C)) — one pass for the table, then about 9 corner sweeps. Space O(R · C).

Solution

"""Sprinkler square — binary search on the block width, checked with 2D prefix sums."""


def prefix_sums(grove):
    """pre[r][c] is the demand of every cell strictly above and left of (r, c)."""
    rows, cols = len(grove), len(grove[0])
    pre = [[0] * (cols + 1) for _ in range(rows + 1)]
    for r in range(rows):
        for c in range(cols):
            pre[r + 1][c + 1] = (
                grove[r][c] + pre[r][c + 1] + pre[r + 1][c] - pre[r][c]
            )
    return pre


def fits(pre, rows, cols, side, budget):
    """True when some side-by-side block costs at most budget."""
    for r in range(rows - side + 1):
        for c in range(cols - side + 1):
            total = (
                pre[r + side][c + side]
                - pre[r][c + side]
                - pre[r + side][c]
                + pre[r][c]          # the corner subtracted twice, added back once
            )
            if total <= budget:
                return True
    return False


def solve(grove, budget):
    rows, cols = len(grove), len(grove[0])
    pre = prefix_sums(grove)
    # Demands are non-negative, so viability is monotone: T T T F F in the side.
    lo, hi = 0, min(rows, cols)
    while lo < hi:                        # invariant: the answer lies in [lo, hi]
        mid = (lo + hi + 1) // 2          # round up, or lo = mid stalls the loop
        if fits(pre, rows, cols, mid, budget):
            lo = mid                      # mid is viable; look for something wider
        else:
            hi = mid - 1                  # mid is already too wide
    return lo


GROVE = [
    [1, 1, 3, 2],
    [4, 2, 1, 1],
    [3, 2, 2, 1],
    [1, 1, 1, 1],
]

TESTS = [
    ((GROVE, 8), 2),
    ((GROVE, 3), 1),
    ((GROVE, 0), 0),
    ((GROVE, 1000), 4),
    (([[1, 1, 1], [1, 1, 1], [1, 1, 1]], 9), 3),
    (([[5]], 4), 0),
    (([[5]], 5), 1),
    (([[0, 0], [0, 0]], 0), 2),
]

Pitfalls

• Rounding the midpoint down while keeping lo = mid. With lo = 2 and hi = 3, mid is 2 again and the loop never ends; a last-true search needs mid = (lo + hi + 1) // 2.
• Starting the range at 1. When budget is under every tree the answer is 0, which [1, min(R, C)] cannot express — you return 1 for a pump that waters nothing.
• Subtracting both strips without adding the shared corner back. Block totals come out low, so the check passes zones the pump cannot feed.
• Letting a corner run off the edge. The last corner for side k is row R - k, column C - k; looping to R reads clipped blocks and accepts them.

Variants

• Kiln shelf size — the same shape with a one-line check instead of a grid sweep.
• Staircase price board — another grid, walked by its ordering rather than halved.