「青训营 X 码上掘金」攒青豆
当青训营遇上码上掘金
题目回顾
现有 n 个宽度为 1 的柱子，给出 n 个非负整数依次表示柱子的高度，排列后如下图所示，此时均匀从上空向下撒青豆，计算按此排列的柱子能接住多少青豆。（不考虑边角堆积）
输入：height = [5,0,2,1,4,0,1,0,3]
输出：17
解析：上面是由数组 [5,0,2,1,4,0,1,0,3] 表示的柱子高度，在这种情况下，可以接 17 个单位的青豆。
解题思路
从理想情况到实际情况，我们可以将问题分解为以下几个步骤：
计算边界柱子之间的青豆数
计算边界柱子之间的青豆数，即计算边界柱子之间的面积减去边界柱子之间的柱子高度。
例题中，边界柱子为 5 和 4，面积为 4*3-3=9，另外边界柱子为 4 和 3，面积为 3*3-1=8，所以总面积为 17。
def calc_area(arr: list) -> int:
'''计算边界内的面积'''
height = min(arr[0], arr[-1])
width = len(arr)-2
s = functools.reduce(lambda x, y: x+y, arr[1:-1])
return height*width-s
查找边界柱子
边界包括数组第一个元素 L 和最后一个元素 R。中间的元素要从 low=min(L, R) 开始，向另一侧遍历，如果遇到比 low 大的元素，则将该元素作为 low ，将索引边界添加到 index 数组中，继续向另一侧遍历。
返回升序的索引数组，方便后续的计算
def find_border(arr: list) -> list:
"""找到边界, 返回边界的索引"""
l = len(arr)
flag = arr[0] < arr[-1]
low = min(arr[0], arr[-1])
index = [0 if flag else l-1]
# 从小的那一边开始遍历
step = 1 if flag else -1
start = 1 if flag else l-2
end = l-2 if flag else 1
for i in range(start, end, step):
if arr[i] > low:
index.append(i)
low = arr[i]
index.append(l-1 if flag else 0)
return sorted(index)
find_border 单元测试
import unittest
from app import find_border
class TestCalculator(unittest.TestCase):
def test_LM(self):
print("find_border---左中测试")
arr = [5, 0, 2, 1, 4, 0, 1, 0, 3]
result = find_border(arr)
self.assertEqual(result, [8, 4, 0])
def test_RM(self):
print("find_border---右中测试")
arr = [3, 0, 1, 0, 4, 1, 2, 0, 5]
result = find_border(arr)
self.assertEqual(result, [0, 4, 8])
def test_L(self):
print("find_border---左测试")
arr = [5, 0, 2, 1, 3, 0, 1, 0, 4]
result = find_border(arr)
self.assertEqual(result, [8, 0])
def test_R(self):
print("find_border---右测试")
arr = [4, 0, 1, 0, 3, 1, 2, 0, 5]
result = find_border(arr)
self.assertEqual(result, [0, 8])
考虑非正常输入
当头尾的柱子为 0 的时候，边界计算就会有误，所以需要对输入数组进行预处理。
def arr_trim(arr: list) -> list:
'''去除数组两端的0'''
l = 0
r = len(arr)-1
while arr[l] == 0 and l < r:
l += 1
while arr[r] == 0 and l < r:
r -= 1
return arr[l:r+1]
arr_trim 单元测试
import unittest
from app import arr_trim
class TestCalculator(unittest.TestCase):
def test_all_zero(self):
print("arr_trim---全零测试")
arr = [0, 0, 0, 0, 0, 0, 0]
result = arr_trim(arr)
self.assertEqual(result, [0])
def test_left_zero(self):
print("arr_trim---左零测试")
arr = [0, 0, 0, 2, 0, 0, 1]
result = arr_trim(arr)
self.assertEqual(result, [2, 0, 0, 1])
def test_right_zero(self):
print("arr_trim---右零测试")
arr = [2, 0, 0, 1, 0, 0, 0]
result = arr_trim(arr)
self.assertEqual(result, [2, 0, 0, 1])
def test_none_zero(self):
print("arr_trim---无零测试")
arr = [2, 0, 0, 1]
result = arr_trim(arr)
self.assertEqual(result, [2, 0, 0, 1])
def test_left_right_zero(self):
print("arr_trim---左右零测试")
arr = [0, 0, 0, 2, 0, 0, 1, 0, 0, 0]
result = arr_trim(arr)
self.assertEqual(result, [2, 0, 0, 1])
主函数
- 对输入进行预处理
- 并得出边界数组
- 对不同边界分别计算青豆数，并求和
if __name__ == "__main__":
height = [5, 0, 2, 1, 4, 0, 1, 0, 3]
arr = arr_trim(height)
border = find_border(arr)
count = sum(calc_area(arr[x:y+1]) for x, y in zip(border[:-1], border[1:]))
print(count)
使用 zip 可以方便同时获得相邻的两个索引
总结
在设计算法时最好从最小单元，最理想的情况开始处理，逐渐增加问题的难度。
