一道算法题：水流过几个格子

作者：DeppWang、原文地址

题目来源

一道做智能的对话机器人公司的算法测试题

题目描述

有一个矩形的水槽被分为了 N * M 个格子，每个格子内的高度都不同，每个格子和直接相邻的格子相通（对角相邻的不相通）。从最左上角的格子注入水，水会向相通的较低的或同样高的格子流动，但不会流向较高的格子。请写程序计算水一共会流经多少个格子？

例如如果格子高度如下分布：

3 | 5 | 1
---------
2 | 1 | 5	  
---------
4 | 2 | 1

则水会从左上角流经 3、2、1 三个格子，答案为 3。

解题方案

思路

需要利用广度优先搜索（队列）的思路
将元素位置作为队列元素，而不使用元素值作为队列元素
将 [0, 0] 加入队列，依次遍历队列，将四周小于当前数值的元素加入队列，直到遍历完

3 | 5 | 1
---------
2 | 1 | 5	  
---------
4 | 2 | 1

   3
[[0,0]]

   3     2
[[0,0],[1,0]]
   ↑     
   3     2     1
[[0,0],[1,0],[1,1]]
         ↑

5 | 5 | 1
---------
2 | 1 | 5	  
---------
4 | 2 | 1

   5
[[0,0]]
   
   5     5    2
[[0,0],[0,1],[1,0]]
   ↑
   5     5     2     1     1
[[0,0],[0,1],[1,0],[0,2],[1,1]]
         ↑
   5     5     2     1     1     
[[0,0],[0,1],[1,0],[0,2],[1,1]]  [1,1] 已存在，跳过
               ↑
   5     5     2     1     1     
[[0,0],[0,1],[1,0],[0,2],[1,1]]
                     ↑
   5     5     2     1     1     
[[0,0],[0,1],[1,0],[0,2],[1,1]]
                           ↑

时间复杂度：O(N^2)。N 为格子数量，因需要判断元素在队列中是否已经存在，所以需要跟其他元素比较
空间复杂度：O(N)。N 为格子数量，最长需要长度为 N 的额外队列

代码

class Solution:
    def printGridNumber(self, height_arr):
        def judge_contain(key, queue):
            for x in queue:
                if x == key:
                    return True
            return False

        def append(queue, i, j):
            xlen = len(height_arr)
            ylen = len(height_arr[0])
            if i - 1 >= 0 and height_arr[i - 1][j] <= height_arr[i][j] and not judge_contain([i - 1, j], queue):
                queue.append([i - 1, j])

            if j - 1 >= 0 and height_arr[i][j - 1] <= height_arr[i][j] and not judge_contain([i, j - 1], queue):
                queue.append([i, j - 1])

            if i + 1 < xlen and height_arr[i + 1][j] <= height_arr[i][j] and not judge_contain([i + 1, j], queue):
                queue.append([i + 1, j])

            if j + 1 < ylen and height_arr[i][j + 1] <= height_arr[i][j] and not judge_contain([i, j + 1], queue):
                queue.append([i, j + 1])

        queue = [[0, 0]]
        append(queue, 0, 0)
        for x in queue:
            append(queue, x[0], x[1])
            # print(queue)
        return len(queue)


if __name__ == '__main__':
    solution = Solution()
    height_arr = [[3, 2, 1], [3, 2, 1], [2, 3, 1]]
    print(solution.printGridNumber(height_arr))

测试用例

[[3, 5, 1], [2, 1, 5], [4, 2, 1]]  # input
3                                  # expect value

[[5, 5, 1], [2, 1, 5], [4, 2, 1]]
5

[[3, 2, 1], [3, 2, 1], [3, 2, 1]]
9

[[3, 2, 1], [3, 2, 1], [2, 3, 1]]
8

[[6, 3, 4, 2], [5, 4, 4, 1], [3, 3, 2, 1]]
12

错误思路

直接遍历所有元素，如果当前值小于左上角、小于周围的值、周围的值小于左上角，则将结果加 1

class Solution:

    def printGridNumber(self, height_arr):
        def judge(i, j):
            left_top_height = height_arr[0][0]
            xlen = len(height_arr)
            ylen = len(height_arr[0])
            if i - 1 >= 0 and left_top_height >= height_arr[i - 1][j] >= height_arr[i][j]:
                return True

            if j - 1 >= 0 and left_top_height >= height_arr[i][j - 1] >= height_arr[i][j]:
                return True

            if i + 1 < xlen and left_top_height >= height_arr[i + 1][j] >= height_arr[i][j]:
                return True

            if j + 1 < ylen and left_top_height >= height_arr[i][j + 1] >= height_arr[i][j]:
                return True

            return False

        if len(height_arr) == 0:
            return
        if len(height_arr[0]) == 0:
            return
        left_top_height = height_arr[0][0]
        # print(left_top_height)
        number = 1
        for i in range(len(height_arr)):
            for j in range(len(height_arr[0])):
                if i == 0 and j == 0:
                    continue
                if height_arr[i][j] <= left_top_height and judge(i, j):
                    number += 1

        return number


if __name__ == '__main__':
    solution = Solution()
    height_arr = [[3,5,1],[2,1,5],[4,2,1]]
    print(solution.printGridNumber(height_arr))