【程序员代码面试指南】：Python：CD1 在行列都排好序的矩阵中找指定的数

描述

给定一个N×M的整形矩阵matrix和一个整数K, matrix的每一行和每一列都是排好序的。
实现一个函数，判断K是否在matrix中

[要求]

时间复杂度为O(N+M)，额外空间复杂度为O(1)。

输入描述：

第一行有三个整数N, M, K
接下来N行，每行M个整数为输入的矩阵

输出描述：

若K存在于矩阵中输出"Yes"，否则输出"No"

示例1

输入：
2 4 5
1 2 3 4
2 4 5 6

输出：
Yes

示例2

输入：
2 4 233
1 2 3 4
2 4 5 6

输出：
No

代码

def func(matrix, M, N, K):
    i, j = 0, M-1
    while i < N and j > 0:
        
        if matrix[i][j] > K:
            j -= 1
        elif matrix[i][j] < K:
            i += 1
        else:
            return "Yes"
    return "No"


N,M,K = map(int, input().split())
matrix = []
for _ in range(N):
    matrix.append(list(map(int, input().split())))
print(func(matrix, M, N, K))