题解 | MOE Top‑k 路由

n,m,p,k = map(int,input().split())
if n % m != 0 or p > m or p*n//m < k:
    print("error")
else:
    arr = list(map(float,input().split()))
    g = int(n // m)
    group = {i:[] for i in range(m)} # 记录概率
    node = {i:[] for i in range(m)} # 记录编号
    for i in range(n):
        index = i // g
        group[index].append(arr[i])
        node[index].append(i)
    for key in group.keys():
        group[key] = max(group[key])
    sort_group = sorted(group.items(),key=lambda x:[-x[1],x[0]])
    sort_group_p = sort_group[:p] 
    group_id = {i[0] for i in sort_group_p} # 获取保留的组号
    zj_id = []
    for i in range(m):
        if i in group_id:
            zj_id += node[i] # 提取保留组的专家编号
    final_arr = sorted({i:arr[i] for i in zj_id}.items(),key=lambda x:[-x[1],x[0]])[:k]
    final_zj_id = sorted([i[0] for i in final_arr])
    print(" ".join([str(i) for i in final_zj_id]))