Search code examples
algorithmdivide-and-conquer

How many rectangles contain exactly k ones on a grid N*M


You are given a grid of 0's and 1's and it's dimension 1 ≤ N,M ≤ 2500 and a number 0 ≤ K ≤ 6. The task is to count the number of rectangles in the grid that have exactly K ones inside it.

It has to be quicker that O(N^2*M), something like O(NMlog(N+M)) will work. My best aproach was a dp with complexity O(N^2*M), but this is to slow, I've been told that the answer is divide and conquer but I can't get it. Any idea?


Solution

  • One way to get the log factor is to divide horizontally, add the number of rectangles with K 1s above the dividing line, the number of rectangles with K 1s below the dividing line (computed recursively), and for each horizontal width, (there are O(|columns|^2) of them) the number of rectangles that extend some part above and some part below the divided line with the same fixed width. We compute the last one by splitting into two-part partitions of K (maxed at 7 since K ≤ 6 and we may want one part zero). We can binary search on a fixed width and bottom or top line, up or down, on matrix prefix sums, which we can precalculate in O(M * N) and retrieve in O(1).

    f(vertical_range) =
      f(top_part) +
      f(bottom_part) +
      NumTopRecs(w, i) * NumBottomRecs(w, k - i)
        where i <- [0...k],
              w <- all O(M^2) widths
    

    The trick is that in each recursive call, we rotate the half that's passed to f such that the horizontal divider line becomes vertical, which means we've cut the M for our call (from which we draw O(M^2) widths) in 2.