Water Flow | InterviewBit | Solution Explained

I explain the solution to Water Flow on InterviewBit in detail, with visuals and diagrams to help understand things better. Code in Python3.

0. Water Flow: Problem Discussion

0.0. IO Format

• input:
  • A: a 2D integer matrix, where each cell represents the height
• output:
  • the number of cells from where water can reach both blue and red lakes
• condition:
  • when water is placed on any cell, it can move to other neighbourhood cells, if they are equal or lesser in height (basically, water can not climb up!)

0.1. Examples

In the case below, we are given the A matrix, with blue and red lakes marked. The cells in green are the places from where we can reach both the lakes, following the condition mentioned in the question. Thus, we return 7 as the answer, the count of the green cells.

1. Water Flow: Observations & Logic

1.0. Intuitions

The first thing that is immediately clear is that we need to think about this grid as a graph. Now, we are looking at possibly implementing DFS or BFS, since we need to traverse the graph – to know if we can reach the lakes or not.

1.1. Inverting the Problem

Here’s the thing. If we start from each cell in the matrix of size M*N, we will have to do a dfs/bfs O(MN) times – and each of them will take O(MN) time again – since we need to traverse (possibly) the entire graph for each cell.

This means that the time complexity is O(M^2 * N^2) well beyond the constraints mentioned in the question.

So, what do we do? It seems like we are stuck in a dead end.

In questions like these, where there seems to be no way out, try inverting the problem on its own head.

That is, instead of starting from one cell and seeing what lake(s) we end up at by flowing down, let’s start from the lakes instead – and let the water climb up!

1.2. Multi-source BFS

Again, as per our choice, we can go for a DFS or a BFS. In my case, I went for a BFS (to avoid recursion limiting issues in Python3).

We start a multi-source BFS from all the cells adjacent to the blue lake – since we know all of these cells will definitely flow into the blue lake. In the image below, I have marked the blue starting cells, the place we start from.

And now, its just a matter of running the multi-source BFS keeping in mind that we are ascending the heights – not descending anymore. Finally, the graph looks as shown below.

Now, we do the same for the red lake; starting a multi-source BFS from the cells adjacent to the lake. The below is the implementation

2. Water Flow: Optimized Implementation

2.0. Code

class Solution:
    def solve(self, A):
        def bfs(curr):
            visited = set()
            while curr:
                temp = []
                for i, j in curr:
                    if (i, j) in visited: continue
                    visited.add((i, j))

                    colors[i][j] += 1
                    for ii, jj in [(i+1, j), (i-1, j), (i, j+1), (i, j-1)]:
                        if not (0 <= ii < m and 0 <= jj < n): continue
                        if A[i][j] <= A[ii][jj]: 
                            temp.append((ii, jj))

                curr = temp
            
            return
        
        m, n = len(A), len(A[0]) # dims of matrix
        # stores the count of lakes it can reach
        colors = [[0]*n for _ in range(m)]

        # stores the sources for the BFS
        curr = []
        for i in range(m): curr.append((i, 0))
        for j in range(n): curr.append((0, j))
        # runs the BFS
        bfs(curr)

        # similarly
        curr = []
        for i in range(m): curr.append((i, n-1))
        for j in range(n): curr.append((m-1, j))
        bfs(curr)

        # count the number of cells in colors that touch 2 lakes
        ans = 0
        for i in range(m):
            for j in range(n):
                if colors[i][j] == 2:
                    ans += 1 
        return ans
        

2.1. Complexity Analysis

• Time: O(MN) to iterate over each cell once in the matrix for each of the blue and red lakes. This gives us O(MN + MN) for blue and red respectively.
• Space: O(MN) to create the colors matrix.

Tanishq Chaudhary

Producing high-quality intuitive explanations of interview problems. Currently covering LeetCode and InterviewBit.