MarsCode-贪心猫的鱼干大分配

贪心猫的鱼干大分配

问题描述

在猫星球上，小R负责给一行排队的猫分发鱼干。每只猫有一个等级，等级越高的猫应该得到更多的鱼干。规则如下：

1. 每只猫至少得到一斤鱼干。
2. 如果一只猫的等级高于它相邻的猫，它就应该得到比相邻的猫更多的鱼干。

小R想知道，为了公平地满足所有猫的等级差异，他至少需要准备多少斤鱼干。

---

测试样例

样例1：

输入：n = 3, cats_levels = [1, 2, 2]
输出：4

样例2：

输入：n = 6, cats_levels = [6, 5, 4, 3, 2, 16]
输出：17

样例3：

输入：n = 20, cats_levels = [1, 2, 2, 3, 3, 20, 1, 2, 3, 3, 2, 1, 5, 6, 6, 5, 5, 7, 7, 4]
输出：35

解答

def solution(n, cats_levels):
    # Please write your code here
    fish = [1 for _ in range(n)] 
    # 从左到右遍历，再从右到左遍历
    for i in range(1, n):
        if cats_levels[i] > cats_levels[i-1]:
            fish[i] = fish[i-1] + 1
    for i in range(n-2, -1, -1):
        if cats_levels[i] > cats_levels[i+1] and fish[i] <= fish[i+1]:
            fish[i] = fish[i+1]+1
    return sum(fish)

if __name__ == "__main__":
    #  You can add more test cases here
    cats_levels1 = [1, 2, 2]
    cats_levels2 = [6, 5, 4, 3, 2, 16]
    cats_levels3 = [1, 2, 2, 3, 3, 20, 1, 2, 3, 3, 2, 1, 5, 6, 6, 5, 5, 7, 7, 4]
    print(solution(3, cats_levels1) == 4)
    print(solution(6, cats_levels2) == 17)
    print(solution(20, cats_levels3) == 35)