378. Kth Smallest Element in a Sorted Matrix

Problem

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.

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 ≤ n2.

Problem Summary

Given an n x n matrix where each row and each column are sorted in ascending order, find the kth smallest element in the matrix. Note that it is the kth smallest element after sorting, not the kth element.

Note: You may assume that the value of k is always valid, 1 ≤ k ≤ n2.

Solution Ideas

• Given a matrix whose rows are sorted and columns are sorted (not sorted by index), find the Kth smallest element in this matrix. Note that the Kth smallest element being sought does not refer to k distinct elements; identical elements may exist.
• The easiest solution to think of is a priority queue. Push the elements in the matrix into the priority queue one by one. Maintain a min-heap, and once there are k elements in the priority queue, the result has been found.
• The optimal solution for this problem is binary search. What is the search space? According to the problem statement, we know that the element in the upper-left corner of the matrix is the smallest, and the element in the lower-right corner is the largest. That is, the first element of the matrix determines the lower bound, and the last element of the matrix determines the upper bound. Binary search over all values in this solution space to find the Kth smallest element. The condition for determining whether it has been found is that the number of elements in the matrix smaller than mid is equal to K. Continuously narrow in on low, so that when low == high, the Kth smallest element has been found. (Because the problem states that the Kth smallest element must exist, when binary search reaches an element, a result will definitely be obtained.)

Code

package leetcode

import (
	"container/heap"
)

// Solution 1: Binary search
func kthSmallest378(matrix [][]int, k int) int {
	m, n, low := len(matrix), len(matrix[0]), matrix[0][0]
	high := matrix[m-1][n-1] + 1
	for low < high {
		mid := low + (high-low)>>1
		// If count is smaller than k, continue binary search in the half with larger values
		if counterKthSmall(m, n, mid, matrix) >= k {
			high = mid
		} else {
			low = mid + 1
		}
	}
	return low
}

func counterKthSmall(m, n, mid int, matrix [][]int) int {
	count, j := 0, n-1
	// Each iteration counts the number of elements smaller than mid
	for i := 0; i < m; i++ {
		// Traverse each row to count the number of elements smaller than mid
		for j >= 0 && mid < matrix[i][j] {
			j--
		}
		count += j + 1
	}
	return count
}

// Solution 2: Priority queue
func kthSmallest3781(matrix [][]int, k int) int {
	if len(matrix) == 0 || len(matrix[0]) == 0 {
		return 0
	}
	pq := &pq{data: make([]interface{}, k)}
	heap.Init(pq)
	for i := 0; i < len(matrix); i++ {
		for j := 0; j < len(matrix[0]); j++ {
			if pq.Len() < k {
				heap.Push(pq, matrix[i][j])
			} else if matrix[i][j] < pq.Head().(int) {
				heap.Pop(pq)
				heap.Push(pq, matrix[i][j])
			} else {
				break
			}
		}
	}
	return heap.Pop(pq).(int)
}

type pq struct {
	data []interface{}
	len  int
}

func (p *pq) Len() int {
	return p.len
}

func (p *pq) Less(a, b int) bool {
	return p.data[a].(int) > p.data[b].(int)
}

func (p *pq) Swap(a, b int) {
	p.data[a], p.data[b] = p.data[b], p.data[a]
}

func (p *pq) Push(o interface{}) {
	p.data[p.len] = o
	p.len++
}

func (p *pq) Head() interface{} {
	return p.data[0]
}

func (p *pq) Pop() interface{} {
	p.len--
	return p.data[p.len]
}