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.

Water Flow: sample test case explanation
Water Flow: sample test case explanation

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.

Water Flow: starting a multi source bfs from the blue cells
Water Flow: starting a multi source bfs from the blue cells

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.

Water Flow: the end of the multi-source bfs
Water Flow: the end of the multi-source bfs

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
        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

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

        # 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.
Avatar for Tanishq Chaudhary

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


    1. Please write the proof.

    2. This solution need an amendment for 2 digit number.

    3. Please write a Proof for this

    Leave a Reply

    Your email address will not be published. Required fields are marked *

    This site uses Akismet to reduce spam. Learn how your comment data is processed.