378. Kth Smallest Element in a Sorted Matrix

发表于 | 更新于: | 分类于 | 本文总阅读量 次
| 字数统计: 3,525 | 阅读时长 ≈ 22
  • 37.8%

https://leetcode.com/problems/kth-smallest-element-in-a-sorted-matrix/#/description

Given a n x n matrix where each of the rows and columns are sorted in ascending order, find the kth smallest element in the matrix.

Note that it is the kth smallest element in the sorted order, not the kth distinct element.

1
2
3
4
5
6
7
8
9
10
Example:

matrix = [
   [ 1,  5,  9],
   [10, 11, 13],
   [12, 13, 15]
],
k = 8,

return 13.

Note:

You may assume k is always valid, 1 ≤ k ≤ n^2.

可以像array一样输入，寻找第k大的元素，转换成array中的top k问题了。但是这种解法效率是O(nlogk)，下面有效率很高的解法，O(#row)的解法，找时间一定好好看看。我现在没有时间看了

49ms, 85.38%, October 18, 2016

https://discuss.leetcode.com/topic/52865/my-solution-using-binary-search-in-c

My solution using Binary Search in C++

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
class Solution {
public:
    int kthSmallest(vector<vector<int>>& matrix, int k)
    {
        int n = matrix.size();
        int le = matrix[0][0], ri = matrix[n - 1][n - 1];
        int mid = 0;
        while (le < ri)
        {
            mid = le + (ri-le)/2;
            int num = 0;
            for (int i = 0; i < n; i++)
            {
                int pos = upper_bound(matrix[i].begin(), matrix[i].end(), mid) - matrix[i].begin();
                num += pos;
            }
            if (num < k)
            {
                le = mid + 1;
            }
            else
            {
                ri = mid;
            }
        }
        return le;
    }
};

https://discuss.leetcode.com/topic/53126/o-n-from-paper-yes-o-rows

O(n) from paper. Yes, O(#rows).

It’s O(n) where n is the number of rows (and columns), not the number of elements. So it’s very efficient. The algorithm is from the paper Selection in X + Y and matrices with sorted rows and columns, which I first saw mentioned by @elmirap (thanks).

The basic idea: Consider the submatrix you get by removing every second row and every second column. This has about a quarter of the elements of the original matrix. And the k-th element (k-th smallest I mean) of the original matrix is roughly the (k/4)-th element of the submatrix. So roughly get the (k/4)-th element of the submatrix and then use that to find the k-th element of the original matrix in O(n) time. It’s recursive, going down to smaller and smaller submatrices until a trivial 2×2 matrix. For more details I suggest checking out the paper, the first half is easy to read and explains things well. Or @zhiqing_xiao’s solution+explanation.

Cool: It uses variants of saddleback search that you might know for example from the Search a 2D Matrix II problem. And it uses the median of medians algorithm for linear-time selection.

Optimization: If k is less than n, we only need to consider the top-left k×k matrix. Similar if k is almost n2. So it’s even O(min(n, k, n^2-k)), I just didn’t mention that in the title because I wanted to keep it simple and because those few very small or very large k are unlikely, most of the time k will be “medium” (and average n2/2).

Implementation: I implemented the submatrix by using an index list through which the actual matrix data gets accessed. If [0, 1, 2, …, n-1] is the index list of the original matrix, then [0, 2, 4, …] is the index list of the submatrix and [0, 4, 8, …] is the index list of the subsubmatrix and so on. This also covers the above optimization by starting with [0, 1, 2, …, k-1] when applicable.

Application: I believe it can be used to easily solve the Find K Pairs with Smallest Sums problem in time O(k) instead of O(k log n), which I think is the best posted so far. I might try that later if nobody beats me to it (if you do, let me know :-). Update: I did that now.

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
class Solution(object):
    def kthSmallest(self, matrix, k):

        # The median-of-medians selection function.
        def pick(a, k):
            if k == 1:
                return min(a)
            groups = (a[i:i+5] for i in range(0, len(a), 5))
            medians = [sorted(group)[len(group) / 2] for group in groups]
            pivot = pick(medians, len(medians) / 2 + 1)
            smaller = [x for x in a if x < pivot]
            if k <= len(smaller):
                return pick(smaller, k)
            k -= len(smaller) + a.count(pivot)
            return pivot if k < 1 else pick([x for x in a if x > pivot], k)

        # Find the k1-th and k2th smallest entries in the submatrix.
        def biselect(index, k1, k2):

            # Provide the submatrix.
            n = len(index)
            def A(i, j):
                return matrix[index[i]][index[j]]
            
            # Base case.
            if n <= 2:
                nums = sorted(A(i, j) for i in range(n) for j in range(n))
                return nums[k1-1], nums[k2-1]

            # Solve the subproblem.
            index_ = index[::2] + index[n-1+n%2:]
            k1_ = (k1 + 2*n) / 4 + 1 if n % 2 else n + 1 + (k1 + 3) / 4
            k2_ = (k2 + 3) / 4
            a, b = biselect(index_, k1_, k2_)

            # Prepare ra_less, rb_more and L with saddleback search variants.
            ra_less = rb_more = 0
            L = []
            jb = n   # jb is the first where A(i, jb) is larger than b.
            ja = n   # ja is the first where A(i, ja) is larger than or equal to a.
            for i in range(n):
                while jb and A(i, jb - 1) > b:
                    jb -= 1
                while ja and A(i, ja - 1) >= a:
                    ja -= 1
                ra_less += ja
                rb_more += n - jb
                L.extend(A(i, j) for j in range(jb, ja))
                
            # Compute and return x and y.
            x = a if ra_less <= k1 - 1 else \
                b if k1 + rb_more - n*n <= 0 else \
                pick(L, k1 + rb_more - n*n)
            y = a if ra_less <= k2 - 1 else \
                b if k2 + rb_more - n*n <= 0 else \
                pick(L, k2 + rb_more - n*n)
            return x, y

        # Set up and run the search.
        n = len(matrix)
        start = max(k - n*n + n-1, 0)
        k -= n*n - (n - start)**2
        return biselect(range(start, min(n, start+k)), k, k)[0]

https://discuss.leetcode.com/topic/52947/c-priority-queue-solution-o-klogn

C++ priority queue solution O(klogn)

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
class Solution {
public:
struct compare
{
    bool operator()(const pair<int,pair<int, int> >& a, const pair<int,pair<int, int> >& b)
    {
        return a.first>b.first;
    }
};
    int kthSmallest(vector<vector<int>>& arr, int k) {
        
        int n=arr.size(),m=arr[0].size();
        
        priority_queue< pair<int,pair<int, int> >, vector<pair<int, pair<int, int> > >, compare > p;
        
        for(int i=0;i<n;i++)
        p.push(make_pair(arr[i][0],make_pair(i,0)));
        
        int x=k,ans;
        while(x--)
        {
            int e=p.top().first;
            int i=p.top().second.first;
            int j=p.top().second.second;
            ans=e;
            p.pop();
            if(j!=m-1)
            p.push(make_pair(arr[i][j+1],make_pair(i,j+1)));
        }
        return ans;
        
    }
};

https://discuss.leetcode.com/topic/54262/c-o-n-time-o-n-space-solution-with-detail-intuitive-explanation

C++ O(n)-time O(n)-space solution with detail intuitive explanation

This thread is inspired by @StefanPochmann’s thread, which mentioned Mirzaian and Arjoandi’s paper.

Preparations

  1. When n==1 (i.e. the matrix is 1x1. n is the number of row), the problem is trival. Hencefore we only consider the case n>=2.
    Rather than finding one k-th element from the matrix, we will select TWO elements (say, k0-th element and k1-th element) simultaneously, such that 0<=k0<=k1 < n * n and k1-k0<=4n. Obviously, if we can complete the aforementioned selection in O(n), we can find the k-th element in O(n) by simply letting k=k0=k1.
  2. Let x0 denote the k0-th element; let x1 denote the k1-th element. Obviously we have x0<=x1.
    Now we will introduce how to select x0 and x1 in O(n).

General idea:

For an nxn matrix, where n is large, we try to select x0 and x1 in a recursive way.

  1. (Determine submatrix) This step constructs one submatrix, whose number of elements will be approximately a quarter of the original matrix. The submatrix is defined as every other row and every other column of the original matrix. The last row and the last column are included too (the reason will be stated in the sequel.) Then the dimension of the matrix is approximately (n/2) x (n/2). The submatrix is recorded by the indices in the original matrix.

Example 1: the original matrix has indices {0, 1, 2, 3, 4}, then the submatrix has indices {0, 2, 4}.

Example 2: the original matrix has indices {0,1, 2, 3, 4, 5}, then the submatrix has indices {0, 2,4, 5}.

  1. (Determine new k’s) This step determines two new k’s (denoted as k0_ and k1_) such that (i) k0_ is the largest possible integer to ensure k0_-th element in the new submatrix (denoted as x0_) is not greater than x0; (ii) k1_ is the smallest possible integer to ensure k1_-th element in the new submatrix (denoted as x1_) is not less than x1. This step is the most tricky step.
1
2
3
k0_ = floor(k0 / 4)
k1_ = floor(k1 / 4) + n + 1 (when n is even)
      floor((k1 + 2 * n + 1) / 4) (when n is odd)

Picture: the way to determine k0_ and k1_: https://drive.google.com/open?id=0By2m48ItFbTeeDFvaS1WcV9qSWM

The picture can also be founded here.

Recall that we mentioned the last row and column shall always be included in the matrix. That is to ensure we can always found the x1_ such that x1_ >= x1.

  1. (Call recursively) Obtainx0_ and x1_ by recursion.
  2. (Partition) Partition all elements in the original nxn elements into three parts: P1={e: e < x0_}, P2={e: x0_ <= e < x1_ }, P3={e: x1_ < e}. We only need to record the cardinality of P1 and P2 (denoted as |P1| and |P2| respectively), and the elements in P2. Obviously, the cardinality of P2 is O(n).
  3. (Get x0 and x1) From the definition of k0_ and k1_, we have |P1| < k0 <= |P1|+|P2|. When |P1| < k0 < |P1|+|P2|, x0 is the k0-|P1|-th element of P2; otherwise x0=x1_. x1 can be determined in a similar way. This action is also O(n).

Complexities:

  • Time: O(n) —– Apply T(n) = T(n/2) + O(n) in the Master’s Theorem.
  • Space: O(n)

C++ Accepted Code:

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
class Solution {
public:
	int kthSmallest(const std::vector<std::vector<int>> & matrix, int k)
	{
		if (k == 1) // guard for 1x1 matrix
		{
			return matrix.front().front();
		}

		size_t n = matrix.size();
		std::vector<size_t> indices(n);
		std::iota(indices.begin(), indices.end(), 0);
		std::array<size_t, 2> ks = { k - 1, k - 1 }; // use zero-based indices
		std::array<int, 2> results = biSelect(matrix, indices, ks);
		return results[0];
	}

private:
	// select two elements from four elements, recursively
	std::array<int, 2> biSelect(
		const std::vector<std::vector<int>> & matrix,
		const std::vector<size_t> & indices,
		const std::array<size_t, 2> & ks)
	// Select both ks[0]-th element and ks[1]-th element in the matrix,
	// where k0 = ks[0] and k1 = ks[1] and n = indices.size() satisfie
	// 0 <= k0 <= k1 < n*n  and  k1 - k0 <= 4n-4 = O(n)   and  n>=2
	{
		size_t n = indices.size();		
		if (n == 2u) // base case of resursion
		{			
			return biSelectNative(matrix, indices, ks);
		}
		
		// update indices
		std::vector<size_t> indices_;
		for (size_t idx = 0; idx < n; idx += 2)
		{
			indices_.push_back(indices[idx]);
		}
		if (n % 2 == 0) // ensure the last indice is included
		{
			indices_.push_back(indices.back());
		}

		// update ks
		// the new interval [xs_[0], xs_[1]] should contain [xs[0], xs[1]]
		// but the length of the new interval should be as small as possible
		// therefore, ks_[0] is the largest possible index to ensure xs_[0] <= xs[0]
		// ks_[1] is the smallest possible index to ensure xs_[1] >= xs[1]
		std::array<size_t, 2> ks_ = { ks[0] / 4, 0 };
		if (n % 2 == 0) // even
		{
			ks_[1] = ks[1] / 4 + n + 1;
		}
		else // odd
		{
			ks_[1] = (ks[1] + 2 * n + 1) / 4;
		}

		// call recursively
		std::array<int, 2> xs_ = biSelect(matrix, indices_, ks_);

		// Now we partipate all elements into three parts:
		// Part 1: {e : e < xs_[0]}.  For this part, we only record its cardinality
		// Part 2: {e : xs_[0] <= e < xs_[1]}. We store the set elementsBetween
		// Part 3: {e : x >= xs_[1]}. No use. Discard.
		std::array<int, 2> numbersOfElementsLessThanX = { 0, 0 };
		std::vector<int> elementsBetween; // [xs_[0], xs_[1])

		std::array<size_t, 2> cols = { n, n }; // column index such that elem >= x
		 // the first column where matrix(r, c) > b
		 // the first column where matrix(r, c) >= a
		for (size_t row = 0; row < n; ++row)
		{
			size_t row_indice = indices[row];
			for (size_t idx : {0, 1})
			{
				while ((cols[idx] > 0)
					&& (matrix[row_indice][indices[cols[idx] - 1]] >= xs_[idx]))
				{
					--cols[idx];
				}
				numbersOfElementsLessThanX[idx] += cols[idx];
			}
			for (size_t col = cols[0]; col < cols[1]; ++col)
			{
				elementsBetween.push_back(matrix[row_indice][indices[col]]);
			}
		}

		std::array<int, 2> xs; // the return value
		for (size_t idx : {0, 1})
		{
			size_t k = ks[idx];
			if (k < numbersOfElementsLessThanX[0]) // in the Part 1
			{
				xs[idx] = xs_[0];
			}
			else if (k < numbersOfElementsLessThanX[1]) // in the Part 2
			{
				size_t offset = k - numbersOfElementsLessThanX[0];
				std::vector<int>::iterator nth = std::next(elementsBetween.begin(), offset);
				std::nth_element(elementsBetween.begin(), nth, elementsBetween.end());
				xs[idx] = (*nth);
			}
			else // in the Part 3
			{
				xs[idx] = xs_[1];
			}
		}
		return xs;
	}

	// select two elements from four elements, using native way
	std::array<int, 2> biSelectNative(
		const std::vector<std::vector<int>> & matrix,
		const std::vector<size_t> & indices,
		const std::array<size_t, 2> & ks)
	{
		std::vector<int> allElements;
		for (size_t r : indices)
		{
			for (size_t c : indices)
			{
				allElements.push_back(matrix[r][c]);
			}
		}
		std::sort(allElements.begin(), allElements.end());
		std::array<int, 2> results;
		for (size_t idx : {0, 1})
		{
			results[idx] = allElements[ks[idx]];
		}
		return results;
	}
};

https://discuss.leetcode.com/topic/52886/c-solution-same-as-find-k-pairs-with-smaller-sums

C++ solution same as Find K pairs with smaller sums

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
class Solution {
public:

    int kthSmallest(vector<vector<int>>& matrix, int k) {
        auto comp = [&matrix](pair<int, int> p1, pair<int, int> p2){
            return matrix[p1.first][p1.second] > matrix[p2.first][p2.second];
        };
        priority_queue<pair<int, int>, vector<pair<int,int>>, decltype(comp)> que(comp);
        que.push(make_pair(0, 0));
        int count = 1;
        while(count < k){
             auto temp = que.top();
             que.pop();
             if(temp.first+1 < matrix.size()){
                 que.push(make_pair(temp.first+1, temp.second));
             }
             if(temp.first == 0 && temp.second+1 < matrix[0].size()){
                 que.push(make_pair(temp.first, temp.second+1));
             }
             count++;
        }
        auto t = que.top();
        return matrix[t.first][t.second];
    }
};

https://discuss.leetcode.com/topic/52866/python-one-line-solution

python one-line solution …

life is too short to figure out more intelligent solution…

1
2
3
4
import heapq
class Solution(object):
    def kthSmallest(self, matrix, k):
        return list(heapq.merge(*matrix))[k-1]

https://discuss.leetcode.com/topic/52912/binary-search-heap-and-sorting-comparison-with-concise-code-and-1-liners-python-72-ms

Binary Search, Heap and Sorting comparison, with concise code and 1-liners, Python 72 ms

For n X n matrix,

  1. Binary search (based on the solution from @光速小子) gives me 72 ms.

The time complexity is O(n log(n) log(N)), where N is the search space that ranges from the smallest element to the biggest element. You can argue that int implies N = 2^32, so log(N) is constant. In a way, this is an O(n * log(n)) solution.

The space complexity is constant.

I thought this idea was weird for a while. Then I noticed the previous problem 377. Combination Sum IV is pretty much doing the same thing, so this idea may actually be intended.

Here is a 8-liner implementation:

1
2
3
4
5
6
7
8
9
10
class Solution(object):
    def kthSmallest(self, matrix, k):
        lo, hi = matrix[0][0], matrix[-1][-1]
        while lo<hi:
            mid = (lo+hi)//2
            if sum(bisect.bisect_right(row, mid) for row in matrix) < k:
                lo = mid+1
            else:
                hi = mid
        return lo
  1. Heap solution gives me 176 ms. The time complexity is O(k log n), so the worst-case and average-case time complexity is O(n^2 log n). Space complexity is O(n).
1
2
3
4
5
6
7
8
9
10
class Solution(object):
    def kthSmallest(self, matrix, k):
        heap = [(row[0], i, 0) for i, row in enumerate(matrix)]
        heapq.heapify(heap)
        ret = 0
        for _ in range(k):
            ret, i, j = heapq.heappop(heap)
            if j+1 < len(matrix[0]):
                heapq.heappush(heap, (matrix[i][j+1], i, j+1))
        return ret
  1. Sorting gives me 80ms. Time complexity of sorting an array of size n^2 is O(n^2 * log n). Space complexity is O(n^2).

The difference is that Timsort implemented in Python is capable of taking advantage of existing partial orderings. Moving sorted data in bulk is always faster than comparing and moving individual data elements, due to modern hardware architecture. Time complexity is the same because merging n sorted arrays of size n is still O(n^2 * log n) in the worst case.

1
2
3
4
5
6
class Solution(object):
    def kthSmallest(self, matrix, k):
        l = []
        for row in matrix:
            l += row
        return sorted(l)[k-1]

Here are some O(n^3) (slow due to list concatenation) 1-liners:

1
2
3
class Solution(object):
    def kthSmallest(self, matrix, k):
        return sorted(sum(matrix, []))[k-1]

and

1
2
3
class Solution(object):
    def kthSmallest(self, matrix, k):
        return sorted(reduce(lambda a,b:a+b, matrix))[k-1]

Here are some O(n^2 * log n) 1-liners provided by @StefanPochmann

1
2
3
return sorted(itertools.chain(*matrix))[k-1]
    return sorted(a for row in matrix for a in row)[k-1]
    return sorted(itertools.chain.from_iterable(matrix))[k-1]

https://discuss.leetcode.com/topic/52862/share-my-python-solution-using-heap

Share My Python Solution using Heap

Algorithm

  1. h: list of tuple (element, row, index), which is initialised with first element of each row in the matrix.
  2. We maintain a heap. In the for loop, we get the smallest element v which is in row r, and replace v with the next element in the row r

Time Complexity

  • insert an element into heap: O(log(n)), where n is the width of the matrix
  • find k the k-th element O(k)
  • Overall: O(klog(n))
    1
    2
    3
    4
    5
    6
    7
    8
    9
    10
    11
    12
    13
    14
    15
    16
    17
    18
    19
    20
    21
    from heapq import heappush, heappop, heapreplace, heapify

    class Solution(object):
        def kthSmallest(self, matrix, k):
            """
            :type matrix: List[List[int]]
            :type k: int
            :rtype: int
            """
            h = [(row[0], row, 1) for row in matrix]
            heapify(h)

            # Since we want to find kth, we pop the first k elements 
            for _ in xrange(k - 1):
                v, r, i = h[0]
                if i < len(r):
                    heapreplace(h, (r[i], r, i + 1))
                else:
                    heappop(h)

            return h[0][0]
谢谢你，可爱的朋友。