在网格中找到第 K 个最小元素？

Question

We are given a grid of size N * N where each element A[i][j] is calculated by this equation (i + j) ^ 2 + (ji) * 10^5 .

We need to find the Kth min element in optimized way.

Constraints :

1 <= number of test cases <= 20 
1 <= N <= 5*10^4 
1 <= k <= N^2

How to solve this problem in an efficient way?

Answer 1

The problem that you are trying to solve (finding kth smallest element in an unordered array) is called the selection problem . There are many ways to solve this problem, and one of the best known ways is the quick select algorithm .

Answer 2

You can solve this problem in O(N log N) time. Note that that is sublinear in the size of the matrix, which is N*N. Of course you should never actually construct the matrix.

First, it helps to see what the values in the matrix looks like :

Note that over the possible range of matrix sizes, the smallest element is always at (i,j) = (N-1,0).

The largest element will be either (0,N-1) or (N-1,N-1), depending on the size of the matrix.

Consider the elements on the line from the smallest to largest. If you pick any one of these elements, then you can trace the contour to find the number of <= elements in O(N) time.

Furthermore, the elements on this line are always monotonically increasing from smallest to largest, so do a binary search on this line of elements, to find the largest one such that the number of <= elements is < k. At O(N) time per test, this takes O(N log N) all together.

Let's say the element you discover has value x, and the number of elements <= x is r, with r < k. Next trace the contour for the next higher element on the line, call its value y, and make list of all the values v such that x < v <= y. There will be O(N) elements in this list, and this takes only O(N) time.

Finally, just use sorting or quickselect to pick the (kr)th element from this list. Again this takes at mostO(N log N) time.