前K小数 码蹄集

题目来源：码蹄集

题目描述：

最直接的思路是先将所有的和计算出来，并进行排序，然后输出前k小的数。

Python代码直接思路（会超内存）：

n, k = map(int, input().split())
a = list(map(int, input().split()))
b = list(map(int, input().split()))

sums = []
for i in range(n):
    for j in range(n):
        sums.append(a[i] + b[j])

sums.sort()
for i in range(k):
    print(sums[i], end=" ")

Python代码优化实现：

import heapq

class Node:
    def __init__(self, sum_, b):
        self.sum = sum_
        self.b = b
        
    def __lt__(self, other):
        return self.sum < other.sum

n, k = map(int, input().split())

a = list(map(int, input().split()))
b = list(map(int, input().split()))

a.sort()
b.sort()

min_list = []
for i in range(n):
    heapq.heappush(min_list, Node(a[i] + b[0], 0))

res = []
for i in range(k):
    node = heapq.heappop(min_list)
    res.append(node.sum)
    temp_b = node.b
    if temp_b + 1 < n:
        heapq.heappush(min_list, Node(node.sum - b[temp_b] + b[temp_b + 1], temp_b + 1))

print(" ".join(str(x) for x in res))

C++代码实现：

代码参考链接：https://betheme.net/dashuju/49396.html?action=onClick

#include<bits/stdc++.h>
using namespace std;

const int maxSize = 10100;
int a[maxSize];
int b[maxSize];

struct Node{
    int sum, b;
    Node(int sum, int b):sum(sum), b(b){}
    bool operator < (const Node &b) const{
        return sum > b.sum;
    }
};

priority_queue<Node> minList;

int main(){
    int n, k;
    cin >> n >> k;

    for(int i = 0; i < n; i++){
        cin >> a[i];
    }

    for(int i = 0; i < n; i++){
        cin >> b[i];
    }

    sort(a, a + n);
    sort(b, b + n);

    for(int i = 0; i < n; i++){
        minList.push(Node(a[i] + b[0], 0));
    }

    for(int i = 0; i < k; i++){
        Node node = minList.top();
        minList.pop();
        cout << node.sum << " ";
        int tempB = node.b;
        if(tempB + 1 < n){
            minList.push(Node(node.sum - b[tempB] + b[tempB + 1], tempB + 1));
        }
    }

    return 0;
}

Java代码实现：

import java.util.*;

public class Main {
    static class Node implements Comparable<Node> {
        int sum;
        int b;

        public Node(int sum, int b) {
            this.sum = sum;
            this.b = b;
        }

        @Override
        public int compareTo(Node other) {
            return Integer.compare(this.sum, other.sum);
        }
    }

    public static void main(String[] args) {
        Scanner sc = new Scanner(System.in);
        int n = sc.nextInt();
        int k = sc.nextInt();

        int[] a = new int[n];
        int[] b = new int[n];

        for (int i = 0; i < n; i++) {
            a[i] = sc.nextInt();
        }

        for (int i = 0; i < n; i++) {
            b[i] = sc.nextInt();
        }

        Arrays.sort(a);
        Arrays.sort(b);

        PriorityQueue<Node> minList = new PriorityQueue<>();

        for (int i = 0; i < n; i++) {
            minList.add(new Node(a[i] + b[0], 0));
        }

        for (int i = 0; i < k; i++) {
            Node node = minList.poll();
            System.out.print(node.sum + " ");
            int tempB = node.b;
            if (tempB + 1 < n) {
                minList.add(new Node(node.sum - b[tempB] + b[tempB + 1], tempB + 1));
            }
        }
    }
}

代码提交测试结果：