-
Notifications
You must be signed in to change notification settings - Fork 18
/
09_Matrix-Search.py
67 lines (56 loc) · 1.52 KB
/
09_Matrix-Search.py
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
#!/usr/bin/python
# coding=utf-8
'''
__author__ = 'sunp'
__date__ = '2019/1/25'
Given an n x m array where all rows and columns are in sorted order, write a function to determine whether the array contains an element x.
contains([[1, 2, 3, 4]
[5, 6, 7, 8]
[9, 10, 11, 12]], 4) = True
'''
def contains1(matrix, x):
# binary search: land the row, then the col
if not matrix or not matrix[0]:
return False
m, n = len(matrix), len(matrix[0])
lasts = [row[-1] for row in matrix]
i = _search(lasts, x)
if i == m:
return False
j = _search(matrix[i], x)
if j == n:
return False
if x == matrix[i][j]:
return True
return False
def _search(arr, x):
lo, hi = 0, len(arr) - 1
while lo <= hi:
mid = (lo + hi) >> 1
if x <= arr[mid]:
hi = mid - 1
else:
lo = mid + 1
return lo
def contains2(matrix, x):
# brute but concise
if not matrix or not matrix[0]:
return False
m, n = len(matrix), len(matrix[0])
i, j = 0, n - 1
while i < m and j >= 0:
if x == matrix[i][j]:
return True
if x < matrix[i][j]:
j -= 1
else:
i += 1
return False
if __name__ == '__main__':
A = [[1, 2, 3, 4], [5, 6, 7, 8], [9, 10, 11, 12]]
for contains in [contains1, contains2]:
assert contains(A, 4)
assert contains(A, 6)
assert contains(A, 9)
assert not contains(A, 0)
assert not contains(A, 13)