AI 刷题 148. 小A的子数组权值题解 | 豆包MarsCode AI刷题239. 小U的早餐选择计算问题描述

239. 小U的早餐选择计算

问题描述

小U同学每天在早餐店有多种主食和饮料可以选择。给定两个数组，staples 表示不同主食的价格，drinks 表示不同饮料的价格。每天早上，小U最多会花费不超过 $x$ 元，他可能选择一份主食、一杯饮料，或只选择主食或饮料中的一个。你的任务是计算小U每天早上有多少种不同的选择，满足其最大花费不超过 $x$ 。

示例

例如，对于输入的主食价格 5 20 5 和饮料价格 5 5 2，如果小U的最大花费为 5，则一共有 5 种不同的选择：

选择主食：选择价格为 5 的主食（2种选择）
选择饮料：选择价格为 5 和 2 的饮料（2种选择）
选择主食和饮料的组合：选择主食 5 和饮料 2，花费为 7，大于最大花费 5，不符合条件。

解题思路

我们可以将问题拆分成几个部分来解决：

单独选择主食或饮料：
- 直接遍历主食和饮料数组，分别计算每个主食和饮料是否符合最大花费 x。
- 选择的数量即为小于等于 x 的主食和饮料的数量之和。
选择主食和饮料的组合：
- 我们可以通过双指针技巧来高效计算主食和饮料的组合。我们首先将主食和饮料数组分别排序。
- 使用两个指针，一个指向主食数组的开始位置，另一个指向饮料数组的结束位置。然后判断当前主食和饮料的组合是否满足不超过 x 的条件。
- 如果组合的价格满足条件，则所有饮料中价格小于当前饮料的组合都可以满足条件，计数时将满足条件的饮料数量加入总数。
- 如果不满足条件，则通过移动饮料指针来尝试减少花费，直到找到符合条件的组合。

解题步骤

对数组排序：首先将 staples 和 drinks 数组分别按升序排序。
双指针法计算组合：使用双指针方法计算可以组合的主食和饮料的数量。
遍历主食和饮料的单独选择：单独遍历 staples 和 drinks，计算各自不超过 x 的数量。
返回最终结果：返回所有满足条件的选择数量。

解题代码

def solution(staples: list, drinks: list, x: int) -> int:
    # 对主食和饮料价格进行排序
    staples.sort()
    drinks.sort()
    
    i, j = 0, len(drinks) - 1
    count = 0
    
    # 使用双指针法计算主食和饮料的组合
    while i < len(staples) and j >= 0:
        if staples[i] + drinks[j] <= x:
            # 如果主食和饮料的组合花费不超过x，则该主食和所有较便宜的饮料组合都符合条件
            count += j + 1
            i += 1
        else:
            # 如果组合的花费超过x，减少饮料的花费
            j -= 1
    
    # 计算只选择主食且花费不超过x的组合数
    for s in staples:
        if s <= x:
            count += 1
    
    # 计算只选择饮料且花费不超过x的组合数
    for d in drinks:
        if d <= x:
            count += 1
    
    return count

# 测试用例
if __name__ == '__main__':
    print(solution(staples=[5, 20, 5], drinks=[5, 5, 2], x=5) == 5)  # 5种选择
    print(solution(staples=[3, 10, 7], drinks=[2, 5], x=7) == 5)      # 5种选择
    print(solution(staples=[10, 15], drinks=[5, 6], x=20) == 7)        # 7种选择

解题思路详细解析

排序：
- 我们首先对主食和饮料数组进行排序，目的是为了使用双指针方法高效地计算符合条件的组合。
- 排序后，主食数组的第一个元素是最便宜的，饮料数组的最后一个元素是最贵的。
双指针法计算组合：
- 使用两个指针 i 和 j 分别指向 staples 和 drinks 数组的开始和结束位置。
- 如果当前主食 staples[i] 和饮料 drinks[j] 的组合不超过 x，那么当前主食与所有较便宜的饮料的组合都符合条件，所以我们可以直接将 j + 1 加到 count 中，并将指针 i 向右移动。
- 如果组合的价格超过了 x，我们就将饮料指针 j 向左移动，尝试减少组合的花费，直到满足条件。
单独选择主食或饮料：
- 遍历 staples 和 drinks 数组，分别计算花费不超过 x 的单独选择数。
返回结果：
- 最终返回计算得到的总选择数。

复杂度分析

时间复杂度：O( $n \log n + m \log m + n + m$ )，其中 $n$ 是主食数组的长度， $m$ 是饮料数组的长度。
- 排序主食数组和饮料数组的时间复杂度是 O( $n \log n$ ) 和 O( $m \log m$ )。
- 双指针法计算主食和饮料组合的时间复杂度是 O( $n + m$ )。
- 计算单独选择的时间复杂度是 O( $n + m$ )。
空间复杂度：O( $1$ )，除了输入和输出数据，算法只使用了常量空间。

AI 刷题 148. 小A的子数组权值 题解 | 豆包MarsCode AI刷题