fork download
  1. #include <iostream>
  2. #include <vector>
  3. #include <queue>
  4. #include <algorithm>
  5. #include <set>
  6. using namespace std;
  7.  
  8. struct Node {
  9. int sum;
  10. int i; // 在数组a中的索引
  11. int j; // 在数组b中的索引
  12.  
  13. Node(int s, int idx_i, int idx_j) : sum(s), i(idx_i), j(idx_j) {}
  14.  
  15. bool operator<(const Node& other) const {
  16. return sum < other.sum; // 最大堆需要重载<
  17. }
  18. };
  19.  
  20. int main() {
  21. int n, k;
  22. cin >> n >> k;
  23.  
  24. vector<int> a(n), b(n);
  25. for (int i = 0; i < n; i++) {
  26. cin >> a[i];
  27. }
  28. for (int i = 0; i < n; i++) {
  29. cin >> b[i];
  30. }
  31.  
  32. // 由于数组是升序的,最大的和肯定是两个数组末尾元素的组合
  33. // 使用最大堆,每次取出当前最大的和
  34. priority_queue<Node> pq;
  35.  
  36. // 使用set来避免重复访问
  37. set<pair<int, int>> visited;
  38.  
  39. // 初始把最大的组合加入堆中
  40. int i = n - 1, j = n - 1;
  41. pq.push(Node(a[i] + b[j], i, j));
  42. visited.insert({i, j});
  43.  
  44. vector<int> result;
  45.  
  46. while (k > 0 && !pq.empty()) {
  47. Node current = pq.top();
  48. pq.pop();
  49.  
  50. result.push_back(current.sum);
  51. k--;
  52.  
  53. if (k == 0) break;
  54.  
  55. // 将可能的候选组合加入堆中
  56. // 候选1: (i-1, j)
  57. if (current.i > 0) {
  58. int new_i = current.i - 1;
  59. int new_j = current.j;
  60. if (visited.find({new_i, new_j}) == visited.end()) {
  61. pq.push(Node(a[new_i] + b[new_j], new_i, new_j));
  62. visited.insert({new_i, new_j});
  63. }
  64. }
  65.  
  66. // 候选2: (i, j-1)
  67. if (current.j > 0) {
  68. int new_i = current.i;
  69. int new_j = current.j - 1;
  70. if (visited.find({new_i, new_j}) == visited.end()) {
  71. pq.push(Node(a[new_i] + b[new_j], new_i, new_j));
  72. visited.insert({new_i, new_j});
  73. }
  74. }
  75. }
  76.  
  77. // 输出结果
  78. for (int i = 0; i < result.size(); i++) {
  79. cout << result[i];
  80. if (i != result.size() - 1) {
  81. cout << " ";
  82. }
  83. }
  84. cout << endl;
  85.  
  86. return 0;
  87. }
Success #stdin #stdout 0.01s 5272KB
stdin
5 4
1 2 3 4 5
3 5 7 9 11
stdout
16 15 14 14