-
Notifications
You must be signed in to change notification settings - Fork 0
/
matrix_search.py
60 lines (49 loc) · 1.53 KB
/
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
"""
Write an efficient algorithm that searches for a value in an m x n matrix.
This matrix has the following properties:
Integers in each row are sorted from left to right.
The first integer of each row is greater than or equal to the last integer of the previous row.
Example:
Consider the following matrix:
[
[1, 3, 5, 7],
[10, 11, 16, 20],
[23, 30, 34, 50]
]
Given target = 3, return 1 ( 1 corresponds to true )
Return 0 / 1 ( 0 if the element is not present, 1 if the element is present ) for this problem
Algorithm does binary search until it finds the right row, then binary search again to find
value in the row. Time: O(log M + log N), Space: O(N)
"""
import numpy as np
def matrix_search(mat, target):
N, M = mat.shape
start_row = 0
end_row = N - 1
while start_row <= end_row:
mid = (start_row + end_row) / 2
mid_row = np.array(mat[mid])[0].tolist()
if mid_row[0] <= target <= mid_row[-1]:
return binary_search(mid_row, target)
elif target > mid_row[-1]:
start_row = mid + 1
elif target < mid_row[0]:
end_row = mid - 1
return 0
def binary_search(arr, target):
start, end = 0, len(arr) - 1
while start <= end:
mid = (start + end) / 2
if arr[mid] == target:
return 1
elif target > arr[mid]:
start = mid + 1
elif target < arr[mid]:
end = mid - 1
return 0
a = np.matrix([
[1, 3, 5, 7],
[10, 11, 16, 20],
[23, 30, 34, 50]
])
print matrix_search(a, 11)