[Google Competitive Programming] Find MEDIAN of for all submatrix size K*K of a matrix

So I encounter this problem that seems easy at first, but then it became very hard.

Given an array int a[n+1][n+1] where 0 <= a[i][j] <= 255 and number K <= n. The array is indexed from (1,1)
Find array f[n+1][n+1] such that:
f(i,j) = median(a[i -> min(i+K-1,n) ][j -> min(j+K-1, n)] ), (the sub matrix from row i->i+K-1, column j->j+K-1).

I.
Here's the basic solution, complexity: O(N^2 * K^2 * log(K^2))

vector<int> elements;
for (int i=1; i<=n; i++)
	for (int j=1; j<=n; j++)
	{
		elements.clear();
		for (int u=i; u<=min(i+K-1, n); u++)
		for (int v=j; v<= min(j+K-1, n); v++)
			elements.append(a[u][v]);
		sort(elements.begin(), elements.end();
		f[i][j] = elements[elements.size() / 2];
	}

II.
The better solution would be using two heaps (or set, etc) and update the mean as we're passing the array.
For example, moving from median(a[1->3][1->3]) to median(a[1->3][2->4]), we only need to remove elements of the first column from the two heap, then add the elements of the 4th column.
-> N^2 * K * log(K^2)

III.
How can I solve this problem in O(N^2) ? Is it possible at all? If anyone can help, I would be very grateful.
Thank you.

Comments (0)