Searching a sorted 2D array


A 2D Array is sorted on its row and columns, find a key.


"""
Program to search an element in a 2-d array
The 2D array has sorted rows and sorted columns
"""
from random import randrange

N=10 # we will generate a NxN array

# initializes a 2-d array sorted on rows and columns
ARRAY = [ [i for i in range(j,N+j)] for j in range(N) ]

KEY = randrange(0,ARRAY[N-1][N-1])

# We will skip the brute force method of O(n^2) 

# Attempt 1 : Use binary search
def binary_search(lst, l, u, key):
    if l > u : return -1;
    m = (l+u)/2
    if lst[m] == key: return m
    if lst[m] < key: return binary_search(lst, m+1,u, key)
    else : return binary_search(lst, l, m-1,key)


def binsearch_method():
    "Use binary search on each row : takes loglinear time: O( n log(n) )"
    for i in range(N):
        idx = binary_search(ARRAY[i],0,len(ARRAY[i])-1,KEY)
        if idx != -1:
            print("found %d at %d,%d" % (KEY, i,idx))
            return

def linear_method():
    "search in linear time"
    currow = 0
    curcol = N-1
    steps = 0
    while curcol > 0 and currow < N:
        steps += 1
        #print(currow, curcol)
        # If the elem in a col is greater than key, we can skip the entire col
        if ARRAY[currow][curcol] > KEY: curcol -= 1 
        # If the first element in a row is > than key or if the last elem is < key, we skip the entire row.
        if ARRAY[currow][0] > KEY or ARRAY[currow][curcol] < KEY: currow += 1
        if ARRAY[currow][curcol] == KEY:
            print("found %d at %d,%d in %d steps" % (KEY, currow,curcol, steps) )
            return

if __name__ == "__main__":
    print("Key is %d" % KEY )
    binsearch_method()
    linear_method()


No comments: