Skip to content

Commit d616e81

Browse files
yanglbmeyanglbme
authored
feat: add solutions to lc problem: No.3478 (#4147)
No.3478.Choose K Elements With Maximum Sum
1 parent 8406aa1 commit d616e81

File tree

7 files changed

+429
-8
lines changed

7 files changed

+429
-8
lines changed

‎solution/3400-3499/3478.Choose K Elements With Maximum Sum/README.md

+148-4
Original file line numberDiff line numberDiff line change
@@ -72,32 +72,176 @@ edit_url: https://github.com/doocs/leetcode/edit/main/solution/3400-3499/3478.Ch
7272

7373
<!-- solution:start -->
7474

75-
### 方法一
75+
### 方法一：排序 + 优先队列（小根堆）
76+
77+
我们可以将数组 $\textit{nums1}$ 转换成一个数组 $\textit{arr}$，其中每个元素是一个二元组 $(x, i)$，表示 $\textit{nums1}[i]$ 的值为 $x$。然后对数组 $\textit{arr}$ 按照 $x$ 进行升序排序。
78+
79+
我们使用一个小根堆 $\textit{pq}$ 来维护数组 $\textit{nums2}$ 中的元素，初始时 $\textit{pq}$ 为空。用一个变量 $\textit{s}$ 来记录 $\textit{pq}$ 中的元素之和。另外，我们用一个指针 $j$ 来维护当前需要添加到 $\textit{pq}$ 中的元素在数组 $\textit{arr}$ 中的位置。
80+
81+
我们遍历数组 $\textit{arr}$，对于第 $h$ 个元素 $(x, i)$，我们将所有满足 $j < h$ 并且 $\textit{arr}[j][0] < x$ 的元素 $\textit{nums2}[\textit{arr}[j][1]]$ 添加到 $\textit{pq}$ 中，并将这些元素的和加到 $\textit{s}$ 中。如果 $\textit{pq}$ 的大小超过了 $k$，我们将 $\textit{pq}$ 中的最小元素弹出，并将其从 $\textit{s}$ 中减去。然后，我们更新 $\textit{ans}[i]$ 的值为 $\textit{s}$。
82+
83+
遍历结束后，返回答案数组 $\textit{ans}$。
84+
85+
时间复杂度 $O(n \times \log n)$，空间复杂度 $O(n)$。其中 $n$ 为数组长度。
7686

7787
<!-- tabs:start -->
7888

7989
#### Python3
8090

8191
```python
82-
92+
class Solution:
93+
def findMaxSum(self, nums1: List[int], nums2: List[int], k: int) -> List[int]:
94+
arr = [(x, i) for i, x in enumerate(nums1)]
95+
arr.sort()
96+
pq = []
97+
s = j = 0
98+
n = len(arr)
99+
ans = [0] * n
100+
for h, (x, i) in enumerate(arr):
101+
while j < h and arr[j][0] < x:
102+
y = nums2[arr[j][1]]
103+
heappush(pq, y)
104+
s += y
105+
if len(pq) > k:
106+
s -= heappop(pq)
107+
j += 1
108+
ans[i] = s
109+
return ans
83110
```
84111

85112
#### Java
86113

87114
```java
88-
115+
class Solution {
116+
public long[] findMaxSum(int[] nums1, int[] nums2, int k) {
117+
int n = nums1.length;
118+
int[][] arr = new int[n][0];
119+
for (int i = 0; i < n; ++i) {
120+
arr[i] = new int[] {nums1[i], i};
121+
}
122+
Arrays.sort(arr, (a, b) -> a[0] - b[0]);
123+
PriorityQueue<Integer> pq = new PriorityQueue<>();
124+
long s = 0;
125+
long[] ans = new long[n];
126+
int j = 0;
127+
for (int h = 0; h < n; ++h) {
128+
int x = arr[h][0], i = arr[h][1];
129+
while (j < h && arr[j][0] < x) {
130+
int y = nums2[arr[j][1]];
131+
pq.offer(y);
132+
s += y;
133+
if (pq.size() > k) {
134+
s -= pq.poll();
135+
}
136+
++j;
137+
}
138+
ans[i] = s;
139+
}
140+
return ans;
141+
}
142+
}
89143
```
90144

91145
#### C++
92146

93147
```cpp
94-
148+
class Solution {
149+
public:
150+
vector<long long> findMaxSum(vector<int>& nums1, vector<int>& nums2, int k) {
151+
int n = nums1.size();
152+
vector<pair<int, int>> arr(n);
153+
for (int i = 0; i < n; ++i) {
154+
arr[i] = {nums1[i], i};
155+
}
156+
ranges::sort(arr);
157+
priority_queue<int, vector<int>, greater<int>> pq;
158+
long long s = 0;
159+
int j = 0;
160+
vector<long long> ans(n);
161+
for (int h = 0; h < n; ++h) {
162+
auto [x, i] = arr[h];
163+
while (j < h && arr[j].first < x) {
164+
int y = nums2[arr[j].second];
165+
pq.push(y);
166+
s += y;
167+
if (pq.size() > k) {
168+
s -= pq.top();
169+
pq.pop();
170+
}
171+
++j;
172+
}
173+
ans[i] = s;
174+
}
175+
return ans;
176+
}
177+
};
95178
```
96179
97180
#### Go
98181
99182
```go
183+
func findMaxSum(nums1 []int, nums2 []int, k int) []int64 {
184+
n := len(nums1)
185+
arr := make([][2]int, n)
186+
for i, x := range nums1 {
187+
arr[i] = [2]int{x, i}
188+
}
189+
ans := make([]int64, n)
190+
sort.Slice(arr, func(i, j int) bool { return arr[i][0] < arr[j][0] })
191+
pq := hp{}
192+
var s int64
193+
j := 0
194+
for h, e := range arr {
195+
x, i := e[0], e[1]
196+
for j < h && arr[j][0] < x {
197+
y := nums2[arr[j][1]]
198+
heap.Push(&pq, y)
199+
s += int64(y)
200+
if pq.Len() > k {
201+
s -= int64(heap.Pop(&pq).(int))
202+
}
203+
j++
204+
}
205+
ans[i] = s
206+
}
207+
return ans
208+
}
209+
210+
type hp struct{ sort.IntSlice }
211+
212+
func (h hp) Less(i, j int) bool { return h.IntSlice[i] < h.IntSlice[j] }
213+
func (h *hp) Push(v any) { h.IntSlice = append(h.IntSlice, v.(int)) }
214+
func (h *hp) Pop() any {
215+
a := h.IntSlice
216+
v := a[len(a)-1]
217+
h.IntSlice = a[:len(a)-1]
218+
return v
219+
}
220+
```
100221

222+
#### TypeScript
223+
224+
```ts
225+
function findMaxSum(nums1: number[], nums2: number[], k: number): number[] {
226+
const n = nums1.length;
227+
const arr = nums1.map((x, i) => [x, i]).sort((a, b) => a[0] - b[0]);
228+
const pq = new MinPriorityQueue();
229+
let [s, j] = [0, 0];
230+
const ans: number[] = Array(k).fill(0);
231+
for (let h = 0; h < n; ++h) {
232+
const [x, i] = arr[h];
233+
while (j < h && arr[j][0] < x) {
234+
const y = nums2[arr[j++][1]];
235+
pq.enqueue(y);
236+
s += y;
237+
if (pq.size() > k) {
238+
s -= pq.dequeue();
239+
}
240+
}
241+
ans[i] = s;
242+
}
243+
return ans;
244+
}
101245
```
102246

103247
<!-- tabs:end -->

‎solution/3400-3499/3478.Choose K Elements With Maximum Sum/README_EN.md

+148-4
Original file line numberDiff line numberDiff line change
@@ -72,32 +72,176 @@ edit_url: https://github.com/doocs/leetcode/edit/main/solution/3400-3499/3478.Ch
7272

7373
<!-- solution:start -->
7474

75-
### Solution 1
75+
### Solution 1: Sorting + Priority Queue (Min-Heap)
76+
77+
We can convert the array $\textit{nums1}$ into an array $\textit{arr}$, where each element is a tuple $(x, i)$, representing the value $x$ at index $i$ in $\textit{nums1}$. Then, we sort the array $\textit{arr}$ in ascending order by $x$.
78+
79+
We use a min-heap $\textit{pq}$ to maintain the elements from the array $\textit{nums2}$. Initially, $\textit{pq}$ is empty. We use a variable $\textit{s}$ to record the sum of the elements in $\textit{pq}$. Additionally, we use a pointer $j$ to maintain the current position in the array $\textit{arr}$ that needs to be added to $\textit{pq}$.
80+
81+
We traverse the array $\textit{arr}$. For the $h$-th element $(x, i)$, we add all elements $\textit{nums2}[\textit{arr}[j][1]]$ to $\textit{pq}$ that satisfy $j < h$ and $\textit{arr}[j][0] < x$, and add these elements to $\textit{s}$. If the size of $\textit{pq}$ exceeds $k$, we pop the smallest element from $\textit{pq}$ and subtract it from $\textit{s}$. Then, we update the value of $\textit{ans}[i]$ to $\textit{s}$.
82+
83+
After traversing, we return the answer array $\textit{ans}$.
84+
85+
The time complexity is $O(n \log n)$, and the space complexity is $O(n)$. Here, $n$ is the length of the array.
7686

7787
<!-- tabs:start -->
7888

7989
#### Python3
8090

8191
```python
82-
92+
class Solution:
93+
def findMaxSum(self, nums1: List[int], nums2: List[int], k: int) -> List[int]:
94+
arr = [(x, i) for i, x in enumerate(nums1)]
95+
arr.sort()
96+
pq = []
97+
s = j = 0
98+
n = len(arr)
99+
ans = [0] * n
100+
for h, (x, i) in enumerate(arr):
101+
while j < h and arr[j][0] < x:
102+
y = nums2[arr[j][1]]
103+
heappush(pq, y)
104+
s += y
105+
if len(pq) > k:
106+
s -= heappop(pq)
107+
j += 1
108+
ans[i] = s
109+
return ans
83110
```
84111

85112
#### Java
86113

87114
```java
88-
115+
class Solution {
116+
public long[] findMaxSum(int[] nums1, int[] nums2, int k) {
117+
int n = nums1.length;
118+
int[][] arr = new int[n][0];
119+
for (int i = 0; i < n; ++i) {
120+
arr[i] = new int[] {nums1[i], i};
121+
}
122+
Arrays.sort(arr, (a, b) -> a[0] - b[0]);
123+
PriorityQueue<Integer> pq = new PriorityQueue<>();
124+
long s = 0;
125+
long[] ans = new long[n];
126+
int j = 0;
127+
for (int h = 0; h < n; ++h) {
128+
int x = arr[h][0], i = arr[h][1];
129+
while (j < h && arr[j][0] < x) {
130+
int y = nums2[arr[j][1]];
131+
pq.offer(y);
132+
s += y;
133+
if (pq.size() > k) {
134+
s -= pq.poll();
135+
}
136+
++j;
137+
}
138+
ans[i] = s;
139+
}
140+
return ans;
141+
}
142+
}
89143
```
90144

91145
#### C++
92146

93147
```cpp
94-
148+
class Solution {
149+
public:
150+
vector<long long> findMaxSum(vector<int>& nums1, vector<int>& nums2, int k) {
151+
int n = nums1.size();
152+
vector<pair<int, int>> arr(n);
153+
for (int i = 0; i < n; ++i) {
154+
arr[i] = {nums1[i], i};
155+
}
156+
ranges::sort(arr);
157+
priority_queue<int, vector<int>, greater<int>> pq;
158+
long long s = 0;
159+
int j = 0;
160+
vector<long long> ans(n);
161+
for (int h = 0; h < n; ++h) {
162+
auto [x, i] = arr[h];
163+
while (j < h && arr[j].first < x) {
164+
int y = nums2[arr[j].second];
165+
pq.push(y);
166+
s += y;
167+
if (pq.size() > k) {
168+
s -= pq.top();
169+
pq.pop();
170+
}
171+
++j;
172+
}
173+
ans[i] = s;
174+
}
175+
return ans;
176+
}
177+
};
95178
```
96179
97180
#### Go
98181
99182
```go
183+
func findMaxSum(nums1 []int, nums2 []int, k int) []int64 {
184+
n := len(nums1)
185+
arr := make([][2]int, n)
186+
for i, x := range nums1 {
187+
arr[i] = [2]int{x, i}
188+
}
189+
ans := make([]int64, n)
190+
sort.Slice(arr, func(i, j int) bool { return arr[i][0] < arr[j][0] })
191+
pq := hp{}
192+
var s int64
193+
j := 0
194+
for h, e := range arr {
195+
x, i := e[0], e[1]
196+
for j < h && arr[j][0] < x {
197+
y := nums2[arr[j][1]]
198+
heap.Push(&pq, y)
199+
s += int64(y)
200+
if pq.Len() > k {
201+
s -= int64(heap.Pop(&pq).(int))
202+
}
203+
j++
204+
}
205+
ans[i] = s
206+
}
207+
return ans
208+
}
209+
210+
type hp struct{ sort.IntSlice }
211+
212+
func (h hp) Less(i, j int) bool { return h.IntSlice[i] < h.IntSlice[j] }
213+
func (h *hp) Push(v any) { h.IntSlice = append(h.IntSlice, v.(int)) }
214+
func (h *hp) Pop() any {
215+
a := h.IntSlice
216+
v := a[len(a)-1]
217+
h.IntSlice = a[:len(a)-1]
218+
return v
219+
}
220+
```
100221

222+
#### TypeScript
223+
224+
```ts
225+
function findMaxSum(nums1: number[], nums2: number[], k: number): number[] {
226+
const n = nums1.length;
227+
const arr = nums1.map((x, i) => [x, i]).sort((a, b) => a[0] - b[0]);
228+
const pq = new MinPriorityQueue();
229+
let [s, j] = [0, 0];
230+
const ans: number[] = Array(k).fill(0);
231+
for (let h = 0; h < n; ++h) {
232+
const [x, i] = arr[h];
233+
while (j < h && arr[j][0] < x) {
234+
const y = nums2[arr[j++][1]];
235+
pq.enqueue(y);
236+
s += y;
237+
if (pq.size() > k) {
238+
s -= pq.dequeue();
239+
}
240+
}
241+
ans[i] = s;
242+
}
243+
return ans;
244+
}
101245
```
102246

103247
<!-- tabs:end -->

‎solution/3400-3499/3478.Choose K Elements With Maximum Sum/Solution.cpp

+30
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,30 @@
1+
class Solution {
2+
public:
3+
vector<long long> findMaxSum(vector<int>& nums1, vector<int>& nums2, int k) {
4+
int n = nums1.size();
5+
vector<pair<int, int>> arr(n);
6+
for (int i = 0; i < n; ++i) {
7+
arr[i] = {nums1[i], i};
8+
}
9+
ranges::sort(arr);
10+
priority_queue<int, vector<int>, greater<int>> pq;
11+
long long s = 0;
12+
int j = 0;
13+
vector<long long> ans(n);
14+
for (int h = 0; h < n; ++h) {
15+
auto [x, i] = arr[h];
16+
while (j < h && arr[j].first < x) {
17+
int y = nums2[arr[j].second];
18+
pq.push(y);
19+
s += y;
20+
if (pq.size() > k) {
21+
s -= pq.top();
22+
pq.pop();
23+
}
24+
++j;
25+
}
26+
ans[i] = s;
27+
}
28+
return ans;
29+
}
30+
};

0 commit comments

Comments
 (0)