【可信模拟】礼物组合

礼物组合

题目描述

某公司员工给安排中秋礼物，有日用品和水果两类别的礼盒，从两类别选择一种形成礼物组合。
现在定义清单，分别表示日用品礼盒和水果礼盒的价格列表。每个清单选择一种，表示一个礼物组合，请计算并输出礼物组合价格比较低的前 num 个：
输出顺序为礼物组合价格从小到大；
若礼物组合价格相同，则把日用品价格低的组合排在前面。

输入

第一个参数是整数数组 goodsPrice，表示每种日用品礼盒的价格，1 <= goodsPrice[i] <= 1000000 且值各不相同
第一个参数是整数数组 fruitsPrice，表示每种水果礼盒的价格，1 <= fruitsPrice[i] <= 1000000 且值各不相同
第三个参数是整数 num，表示需要输出的价格较低的礼物组合的个数，1 <= num <= goodsSize * fruitsSize 且 num <= 10000

输出

礼物组合，每个元素表示 [price1, price2] 分别为日用品礼盒和水果礼盒的价格。

样例1

输入

[100, 150, 125, 130, 110]
[50, 60, 65, 45]
3

输出

[[100, 45], [100, 50], [110, 45]]

解释：共有 20 种礼物组合，价格较低的前 3 种为：
[100, 45]、[100, 50] 和 [110, 45]。

样例2

输入：

[170, 100, 140, 150, 130, 125, 120, 110]
[55, 50, 80, 60, 65, 45]
10

输出：

[[100, 45], [100, 50], [100, 55], [110, 45], [110, 60], [120, 45], [110, 60]]

解释：按礼物组合价格从小到大选择 10 个组合，其中对于组合价格相同的：[100, 55] 和 [110, 45]，要把日用品价格小的排在前面，因此 [100, 55] 排在前面。

solution.py

from typing import List, Tuple
import heapq

class Solution:
    def get_top_t_pack(self, goods_price: List[int], fruits_price: List[int], num: int) -> List[Tuple[int, int]]:
        min_heap = []
        fruits_price.sort()
        # 生成初始组合，将每个日用品与第一个水果组合
        for g_price in goods_price:
            # 将组合的总价格、日用品价格和水果价格放入堆中
            heapq.heappush(min_heap, (g_price + fruits_price[0], g_price, fruits_price[0], 0))

        result = []

        # 取出前 num 个组合
        while num > 0 and min_heap:
            total_price, g_price, f_price, f_index = heapq.heappop(min_heap)  # 取出最小的组合
            result.append([g_price, f_price])  # 记录组合

            # 如果当前水果组合不是最后一个水果，生成下一个组合
            if f_index + 1 < len(fruits_price):
                new_f_price = fruits_price[f_index + 1]
                heapq.heappush(min_heap, (g_price + new_f_price, g_price, new_f_price, f_index + 1))

            num -= 1

        return result

main.py

import sys
import json
import solution


# 以下为考题输入输出框架，此部分代码不建议改动；提交代码时请勿包含下面代码。
def get_input_lines():
    text = ""
    while True:
        para = sys.stdin.readline().strip()
        if not para:
            break
        text += f"{para}\n"
    lines = text.replace(",\n", ",").splitlines()  # 如果存在换行输入的矩阵数据，将去除换行符
    return lines


def paras_check(paras, paras_default):
    def check(para, dft):
        if isinstance(dft, int) or isinstance(dft, bool) or isinstance(dft, str) or isinstance(dft, float):
            if not isinstance(para, type(dft)):
                return False, f'expect type: {str(type(dft))}, but got value: {json.dumps(para)}'
            return True, ''

        for i in range(len(para)):
            dft_check = dft[0] if isinstance(dft, list) else dft[i]
            res, msg = check(para[i], dft_check)
            if not res:
                return False, f'{json.dumps(para[i])} - {msg}'
            para[i] = tuple(x for x in para[i]) if isinstance(dft_check, tuple) else para[i]
        return True, ''

    return check(paras, paras_default)


def main_process():
    try:
        lines = get_input_lines()
        paras = [json.loads(x) for x in lines]
    except Exception as e:
        print("input format error: " + str(e))
        return

    paras_default = ([0], [0], 0)
    res, msg = paras_check(paras, paras_default)
    if not res:
        print(msg)
        return

    s = solution.Solution()
    result = s.get_top_t_pack(*paras)
    print(json.dumps(result))


main_process()