赛码网 KTV 贪心 排序 / 优先队列

题目

思路

贪心。
首先用nums数组记录想唱每首歌曲的人数，最优策略肯定是选择"最火"的那首歌曲来唱，也就是满足最多人需求的那首歌曲，所以我们对nums数组按人数大到小排序，每次选择一首歌曲，然后更新nums[i]：
nums[i] = nums[i] - min(想唱当前歌曲的人数, x)
直到唱完所有y首歌曲，如果此时nums数组已经全部为0，说明所有人的需求都满足了，输出‘YES’，否则输出‘NO’。

AC代码

（1）排序

import functools

class Song:
    def __init__(self, num=0, song_id=0):
        self.num = num
        self.song_id = song_id

def cmp(song1 ,song2):
    return song2.num - song1.num

T = int(input())

for _ in range(T):
    n, x, y = [int(i) for i in input().split()]
    song_list = []
    # 记录想唱每首歌曲的人数
    nums = [0] * 1005
    
    for _ in range(n):
        a = [int(i) for i in input().split()]
        for i in range(1, len(a)):
            nums[a[i]] += 1
    for i in range(1005):
        if nums[i] != 0:
            song_list.append(Song(nums[i], i))
    
    song_list.sort(key=functools.cmp_to_key(cmp))
    cnt = 0 # 当前唱的歌的数量
    res = True
    while cnt < y:
        song = song_list[0]
        tmp = min(song.num, x)
        nums[song.song_id] -= tmp
        song.num -= tmp
        # 需要重新进行排序
        song_list.sort(key=functools.cmp_to_key(cmp))
        cnt += 1

    for item in song_list:
        if nums[item.song_id] > 0:
            res = False
            break
    print('YES' if res == True else 'NO')

（2）使用优先队列

import heapq

class Song:
    def __init__(self, num=0, song_id=0):
        self.num = num
        self.song_id = song_id

class My_PriorityQueue(object):
    def __init__(self):
        self._queue = []
        self._index = 0

    def push(self, item, priority):
        """
        队列由 (priority, index, item) 形式组成
        priority 增加 "-" 号是因为 heappush 默认是最小堆
        index 是为了当两个对象的优先级一致时，按照插入顺序排列
        """
        heapq.heappush(self._queue, (-priority, self._index, item))
        self._index += 1

    def pop(self):
        """弹出优先级最高的对象"""
        return heapq.heappop(self._queue)[-1]

    def qsize(self):
        return len(self._queue)

    def empty(self):
        return True if not self._queue else False

T = int(input())

for _ in range(T):
    n, x, y = [int(i) for i in input().split()]
    nums = [0] * 1005
    pq = My_PriorityQueue()
    
    for _ in range(n):
        a = [int(i) for i in input().split()]
        for i in range(1, len(a)):
            nums[a[i]] += 1  
    for i in range(1005):
        if nums[i] != 0:
            song = Song(nums[i], i)
            pq.push(song, song.num)
    
    cnt = 0 # 当前唱的歌的数量
    while not pq.empty() and cnt < y:
        song = pq.pop()
        tmp = min(song.num, x)
        nums[song.song_id] -= tmp
        song.num -= tmp
        pq.push(song, song.num)
        cnt += 1
    
    res = True
    for i in range(1, 1005):
        if nums[i] > 0:
            res = False
            break
    print('YES' if res == True else 'NO')

结果