package list

// Problem Definition: https://leetcode.cn/problems/search-a-2d-matrix-ii/

func SearchMatrix(matrix [][]int, target int) bool {
	n, m := len(matrix), len(matrix[0])
	i, j := 0, m-1
	for i < n && j >= 0 {
		switch {
		case target == matrix[i][j]:
			return true
		case target > matrix[i][j]:
			i++
		default:
			j--
		}
	}

	return false
}

// Problem Definition: https://leetcode.cn/problems/kth-smallest-element-in-a-sorted-matrix/

func kthSmallest(matrix [][]int, k int) int {
	n := len(matrix)
	left, right := matrix[0][0], matrix[n-1][n-1]
	for left < right {
		mid := left + (right-left)/2
		if check(matrix, mid, k, n) {
			right = mid
		} else {
			left = mid + 1
		}
	}

	return left
}

func check(matrix [][]int, mid, k, n int) bool {
	i, j := n-1, 0
	num := 0
	for i >= 0 && j < n {
		if matrix[i][j] <= mid {
			num += (i + 1)
			j++
		} else {
			i--
		}
	}

	return num >= k
}