count_lakes google question!

count_lakes.py

import collections

# 0 is water, 1 is land.

arr = [
    [0, 1, 0, 0, 0, 0, 1, 1, 1, 0, 0],
    [0, 0, 0, 0, 0, 0, 1, 0, 1, 0, 0],
    [0, 0, 1, 1, 1, 1, 1, 1, 1, 0, 0],
    [0, 0, 1, 0, 1, 0, 1, 0, 1, 0, 0],
    [0, 0, 1, 0, 1, 0, 1, 0, 1, 0, 0],
    [0, 0, 1, 1, 1, 1, 1, 1, 1, 0, 0],
    [0, 0, 0, 0, 0, 0, 1, 0, 1, 0, 0],
    [0, 0, 0, 0, 0, 0, 1, 1, 1, 0, 0],
    [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
    [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
]

dirs = [(0, -1), (-1, -1), (-1, 0), (-1, 1), (0, 1), (1, 1), (1, 0), (1, -1)]


def fill_sea(image):
    ret = 0
    rows = len(image)
    cols = len(image[0])
    sea = set()
    q = collections.deque([(0, 0)])
    sea.add((0, 0))
    
    # collecting all sea co-ordinates here.
    while q:
        x, y = q.popleft()
        for r, c in dirs:
            nr, nc = x + r, y + c
            if 0 <= nr < rows and 0 <= nc < cols and not image[nr][nc]:
                if (nr, nc) not in sea:
                    sea.add((nr, nc))
                    q.append((nr, nc))

    # count lakes (similar to count islands pretty much)
    for x in range(rows):
        for y in range(cols):
            if not image[x][y] and (x, y) not in sea:
                q = collections.deque([(x, y)])
                sea.add((x, y))
                while q:
                    r, c = q.popleft()
                    for i, j in dirs:
                        nr, nc = r + i, j + c
                        if not image[nr][nc] and (nr, nc) not in sea:
                            sea.add((nr, nc))
                            q.append((nr, nc))
                ret += 1
    return ret


def count_lakes(arr, tup):
    rows = len(arr)
    cols = len(arr[0])
    seen = set()
    seen.add(tup)
    x, y = tup
    q = collections.deque([(x, y)])

    # put all the lands as long as you see them!
    # simple BFFs here.
    while q:
        x, y = q.popleft()
        for r, c in dirs:
            nr, nc = x + r, y + c
            if 0 <= nr < rows and 0 <= nc < cols and arr[nr][nc]:
                if (nr, nc) not in seen:
                    seen.add((nr, nc))
                    q.append((nr, nc))

    # find the min and max now (the imaginary boundary)
    # can do it above as well, not let's do it seperately.
    minrow = rows
    maxrow = maxcol = 0
    mincol = cols

    for r, c in seen:
        minrow = min(minrow, r)
        maxrow = max(maxrow, r)
        mincol = min(mincol, c)
        maxcol = max(maxcol, c)

    # here we are shrinking the image to the lowest rectangle
    # and adding a water boundary around the extracted "smaller" grid.
    image = [
        [0] + arr[i][mincol : maxcol + 1] + [0]
        for i in range(minrow, maxrow + 1)
    ]
    l = len(image[0])
    image = [[0] * l] + image + [[0] * l]

    # reduced grid now looks like ->
    """
    # grid without water added in the boundary (2,2)
    0, 0, 0, 0, 1, 1, 1
    0, 0, 0, 0, 1, 0, 1
    1, 1, 1, 1, 1, 1, 1
    1, 0, 1, 0, 1, 0, 1
    1, 0, 1, 0, 1, 0, 1
    1, 1, 1, 1, 1, 1, 1
    0, 0, 0, 0, 1, 0, 1
    0, 0, 0, 0, 1, 1, 1
    """
    for i in range(len(image)):
        print(image[i])

    # count lakes now...
    print(fill_sea(image))

count_lakes(arr, (2, 2))