巧克力板选择问题 | 豆包MarsCode AI 刷题

小M在春游时打算携带尽可能多的巧克力板。她拥有n块巧克力板，每块巧克力的边长为$a_i$，重量为$a_i^2$。小M有多个不同大小的背包，每个背包都有一定的最大承重限制。她希望你帮助她计算在每个背包的最大承重范围内，最多可以带走多少块巧克力板。

例如：小M有5块巧克力板，边长分别为1, 2, 2, 4, 5，有5个不同的背包，其最大承重分别为1, 3, 7, 9, 15。对于每个背包，她能装入的巧克力块数分别是1, 1, 2, 3, 3。

测试样例

样例1：

输入：n = 5, m = 5, a = [1, 2, 2, 4, 5], queries = [1, 3, 7, 9, 15]
输出：[1, 1, 2, 3, 3]

解析：贪心算法， 直接把a进行排序， 从小到大将巧克力放入背包中，记录相应数量的巧克力需要多少空间，最后判断已有的背包可以满足哪些数量的巧克力

例如：对于样例1，依次从小到大将巧克力放入背包所需的空间是v [1,3,5,9,14]，之后依次遍历queries中背包的大小，对于大小为1，判断其大于v[0]而小于v[1]，因此空间大小为1的背包最多能存放1个巧克力，而对于空间大小为9的背包，正好等于索引为3的值，因此可以放入4个巧克力。按这个规则依次判断即可

具体步骤如下：

1. 初始化测试数据：定义了一个整数列表 a 和一个查询列表 queries。
2. 调用 solution 函数：将 a 和 queries 作为参数传递给 solution 函数，并将返回的结果存储在 result 变量中。
3. 验证结果：使用 == 运算符比较 result 和预期结果 [2, 3, 4, 4] 是否相等。如果相等，打印 True，表示测试通过；如果不相等，打印 False，表示测试失败。

总结来说，这段代码的目的是通过一个具体的例子来验证 solution 函数是否能够正确地计算出每个查询值在整数列表中对应的索引值。

python代码如下

def solution(n: int, m: int, a: list, queries: list) -> list:
    # write code here
    res = []
    info = dict()
    a.sort()
    queries2 = queries.copy()
    queries2.sort()
    heavy = 0
    num = 0
    index = 0
    for item in a:
        heavy += item ** 2
        while index < len(queries) and heavy > queries2[index]:
            info[queries2[index]] = num
            index+=1
        num += 1
    for item in queries:
        res.append(info[item])
    return res

if __name__ == '__main__':
    print(solution(5, 5, [1, 2, 2, 4, 5], [1, 3, 7, 9, 15]) == [1, 1, 2, 3, 3])
    print(solution(4, 3, [3, 1, 2, 5], [5, 10, 20]) == [2, 2, 3])
    print(solution(6, 4, [1, 3, 2, 2, 4, 6], [8, 12, 18, 25]) == [2, 3, 4, 4])

可以看到，我对背包数组进行了排序并用一个字典记录处理，原因是这样可以降低一定的时间复杂度。将背包用index索引从小到大遍历，若循环遍历到的对于数量的巧克力大于当前遍历到的背包空间大小，则记录下对应的巧克力数量，最后将queries的背包大小与字典中的巧克力数量一一对应的返回即可。