Data Structures19 sections · 729 units
Open in Course

Forest Queries Solution

2D prefix sums

function solve(grid, queries):
    n = grid.length
    // Build 2D prefix sum
    prefix = 2D array of size (n+1) x (n+1), filled with 0
    for i = 0 to n - 1:
        for j = 0 to n - 1:
            if grid[i][j] == "*":
                cell = 1
            else:
                cell = 0
            prefix[i+1][j+1] = (prefix[i][j+1] + prefix[i+1][j]
                               - prefix[i][j] + cell)

    results = empty list
    for each (r1, c1, r2, c2) in queries:
        count = (prefix[r2][c2] - prefix[r1-1][c2]
                - prefix[r2][c1-1] + prefix[r1-1][c1-1])
        results.append(count)

    return results

Time: O(nm+q)O(nm + q). The inclusion-exclusion formula handles overlapping corners.