From 323dbe9f7009b91093a5f76672f8f5a18d03baec Mon Sep 17 00:00:00 2001 From: Libin YANG Date: Sun, 9 Mar 2025 21:09:12 +0800 Subject: [PATCH] feat: add solutions to lc problem: No.3478 No.3478.Choose K Elements With Maximum Sum --- .../README.md | 152 +++++++++++++++++- .../README_EN.md | 152 +++++++++++++++++- .../Solution.cpp | 30 ++++ .../Solution.go | 37 +++++ .../Solution.java | 28 ++++ .../Solution.py | 18 +++ .../Solution.ts | 20 +++ 7 files changed, 429 insertions(+), 8 deletions(-) create mode 100644 solution/3400-3499/3478.Choose K Elements With Maximum Sum/Solution.cpp create mode 100644 solution/3400-3499/3478.Choose K Elements With Maximum Sum/Solution.go create mode 100644 solution/3400-3499/3478.Choose K Elements With Maximum Sum/Solution.java create mode 100644 solution/3400-3499/3478.Choose K Elements With Maximum Sum/Solution.py create mode 100644 solution/3400-3499/3478.Choose K Elements With Maximum Sum/Solution.ts diff --git a/solution/3400-3499/3478.Choose K Elements With Maximum Sum/README.md b/solution/3400-3499/3478.Choose K Elements With Maximum Sum/README.md index 57e49332ee980..99fa32054bb64 100644 --- a/solution/3400-3499/3478.Choose K Elements With Maximum Sum/README.md +++ b/solution/3400-3499/3478.Choose K Elements With Maximum Sum/README.md @@ -72,32 +72,176 @@ edit_url: https://github.com/doocs/leetcode/edit/main/solution/3400-3499/3478.Ch -### 方法一 +### 方法一：排序 + 优先队列（小根堆） + +我们可以将数组 $\textit{nums1}$ 转换成一个数组 $\textit{arr}$，其中每个元素是一个二元组 $(x, i)$，表示 $\textit{nums1}[i]$ 的值为 $x$。然后对数组 $\textit{arr}$ 按照 $x$ 进行升序排序。 + +我们使用一个小根堆 $\textit{pq}$ 来维护数组 $\textit{nums2}$ 中的元素，初始时 $\textit{pq}$ 为空。用一个变量 $\textit{s}$ 来记录 $\textit{pq}$ 中的元素之和。另外，我们用一个指针 $j$ 来维护当前需要添加到 $\textit{pq}$ 中的元素在数组 $\textit{arr}$ 中的位置。 + +我们遍历数组 $\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}$。 + +遍历结束后，返回答案数组 $\textit{ans}$。 + +时间复杂度 $O(n \times \log n)$，空间复杂度 $O(n)$。其中 $n$ 为数组长度。 #### Python3 ```python - +class Solution: + def findMaxSum(self, nums1: List[int], nums2: List[int], k: int) -> List[int]: + arr = [(x, i) for i, x in enumerate(nums1)] + arr.sort() + pq = [] + s = j = 0 + n = len(arr) + ans = [0] * n + for h, (x, i) in enumerate(arr): + while j < h and arr[j][0] < x: + y = nums2[arr[j][1]] + heappush(pq, y) + s += y + if len(pq) > k: + s -= heappop(pq) + j += 1 + ans[i] = s + return ans ``` #### Java ```java - +class Solution { + public long[] findMaxSum(int[] nums1, int[] nums2, int k) { + int n = nums1.length; + int[][] arr = new int[n][0]; + for (int i = 0; i < n; ++i) { + arr[i] = new int[] {nums1[i], i}; + } + Arrays.sort(arr, (a, b) -> a[0] - b[0]); + PriorityQueue pq = new PriorityQueue<>(); + long s = 0; + long[] ans = new long[n]; + int j = 0; + for (int h = 0; h < n; ++h) { + int x = arr[h][0], i = arr[h][1]; + while (j < h && arr[j][0] < x) { + int y = nums2[arr[j][1]]; + pq.offer(y); + s += y; + if (pq.size() > k) { + s -= pq.poll(); + } + ++j; + } + ans[i] = s; + } + return ans; + } +} ``` #### C++ ```cpp - +class Solution { +public: + vector findMaxSum(vector& nums1, vector& nums2, int k) { + int n = nums1.size(); + vector> arr(n); + for (int i = 0; i < n; ++i) { + arr[i] = {nums1[i], i}; + } + ranges::sort(arr); + priority_queue, greater> pq; + long long s = 0; + int j = 0; + vector ans(n); + for (int h = 0; h < n; ++h) { + auto [x, i] = arr[h]; + while (j < h && arr[j].first < x) { + int y = nums2[arr[j].second]; + pq.push(y); + s += y; + if (pq.size() > k) { + s -= pq.top(); + pq.pop(); + } + ++j; + } + ans[i] = s; + } + return ans; + } +}; ``` #### Go ```go +func findMaxSum(nums1 []int, nums2 []int, k int) []int64 { + n := len(nums1) + arr := make([][2]int, n) + for i, x := range nums1 { + arr[i] = [2]int{x, i} + } + ans := make([]int64, n) + sort.Slice(arr, func(i, j int) bool { return arr[i][0] < arr[j][0] }) + pq := hp{} + var s int64 + j := 0 + for h, e := range arr { + x, i := e[0], e[1] + for j < h && arr[j][0] < x { + y := nums2[arr[j][1]] + heap.Push(&pq, y) + s += int64(y) + if pq.Len() > k { + s -= int64(heap.Pop(&pq).(int)) + } + j++ + } + ans[i] = s + } + return ans +} + +type hp struct{ sort.IntSlice } + +func (h hp) Less(i, j int) bool { return h.IntSlice[i] < h.IntSlice[j] } +func (h *hp) Push(v any) { h.IntSlice = append(h.IntSlice, v.(int)) } +func (h *hp) Pop() any { + a := h.IntSlice + v := a[len(a)-1] + h.IntSlice = a[:len(a)-1] + return v +} +``` +#### TypeScript + +```ts +function findMaxSum(nums1: number[], nums2: number[], k: number): number[] { + const n = nums1.length; + const arr = nums1.map((x, i) => [x, i]).sort((a, b) => a[0] - b[0]); + const pq = new MinPriorityQueue(); + let [s, j] = [0, 0]; + const ans: number[] = Array(k).fill(0); + for (let h = 0; h < n; ++h) { + const [x, i] = arr[h]; + while (j < h && arr[j][0] < x) { + const y = nums2[arr[j++][1]]; + pq.enqueue(y); + s += y; + if (pq.size() > k) { + s -= pq.dequeue(); + } + } + ans[i] = s; + } + return ans; +} ``` diff --git a/solution/3400-3499/3478.Choose K Elements With Maximum Sum/README_EN.md b/solution/3400-3499/3478.Choose K Elements With Maximum Sum/README_EN.md index e0c6f7167224f..9832f8c87ecfe 100644 --- a/solution/3400-3499/3478.Choose K Elements With Maximum Sum/README_EN.md +++ b/solution/3400-3499/3478.Choose K Elements With Maximum Sum/README_EN.md @@ -72,32 +72,176 @@ edit_url: https://github.com/doocs/leetcode/edit/main/solution/3400-3499/3478.Ch -### Solution 1 +### Solution 1: Sorting + Priority Queue (Min-Heap) + +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$. + +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}$. + +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}$. + +After traversing, we return the answer array $\textit{ans}$. + +The time complexity is $O(n \log n)$, and the space complexity is $O(n)$. Here, $n$ is the length of the array. #### Python3 ```python - +class Solution: + def findMaxSum(self, nums1: List[int], nums2: List[int], k: int) -> List[int]: + arr = [(x, i) for i, x in enumerate(nums1)] + arr.sort() + pq = [] + s = j = 0 + n = len(arr) + ans = [0] * n + for h, (x, i) in enumerate(arr): + while j < h and arr[j][0] < x: + y = nums2[arr[j][1]] + heappush(pq, y) + s += y + if len(pq) > k: + s -= heappop(pq) + j += 1 + ans[i] = s + return ans ``` #### Java ```java - +class Solution { + public long[] findMaxSum(int[] nums1, int[] nums2, int k) { + int n = nums1.length; + int[][] arr = new int[n][0]; + for (int i = 0; i < n; ++i) { + arr[i] = new int[] {nums1[i], i}; + } + Arrays.sort(arr, (a, b) -> a[0] - b[0]); + PriorityQueue pq = new PriorityQueue<>(); + long s = 0; + long[] ans = new long[n]; + int j = 0; + for (int h = 0; h < n; ++h) { + int x = arr[h][0], i = arr[h][1]; + while (j < h && arr[j][0] < x) { + int y = nums2[arr[j][1]]; + pq.offer(y); + s += y; + if (pq.size() > k) { + s -= pq.poll(); + } + ++j; + } + ans[i] = s; + } + return ans; + } +} ``` #### C++ ```cpp - +class Solution { +public: + vector findMaxSum(vector& nums1, vector& nums2, int k) { + int n = nums1.size(); + vector> arr(n); + for (int i = 0; i < n; ++i) { + arr[i] = {nums1[i], i}; + } + ranges::sort(arr); + priority_queue, greater> pq; + long long s = 0; + int j = 0; + vector ans(n); + for (int h = 0; h < n; ++h) { + auto [x, i] = arr[h]; + while (j < h && arr[j].first < x) { + int y = nums2[arr[j].second]; + pq.push(y); + s += y; + if (pq.size() > k) { + s -= pq.top(); + pq.pop(); + } + ++j; + } + ans[i] = s; + } + return ans; + } +}; ``` #### Go ```go +func findMaxSum(nums1 []int, nums2 []int, k int) []int64 { + n := len(nums1) + arr := make([][2]int, n) + for i, x := range nums1 { + arr[i] = [2]int{x, i} + } + ans := make([]int64, n) + sort.Slice(arr, func(i, j int) bool { return arr[i][0] < arr[j][0] }) + pq := hp{} + var s int64 + j := 0 + for h, e := range arr { + x, i := e[0], e[1] + for j < h && arr[j][0] < x { + y := nums2[arr[j][1]] + heap.Push(&pq, y) + s += int64(y) + if pq.Len() > k { + s -= int64(heap.Pop(&pq).(int)) + } + j++ + } + ans[i] = s + } + return ans +} + +type hp struct{ sort.IntSlice } + +func (h hp) Less(i, j int) bool { return h.IntSlice[i] < h.IntSlice[j] } +func (h *hp) Push(v any) { h.IntSlice = append(h.IntSlice, v.(int)) } +func (h *hp) Pop() any { + a := h.IntSlice + v := a[len(a)-1] + h.IntSlice = a[:len(a)-1] + return v +} +``` +#### TypeScript + +```ts +function findMaxSum(nums1: number[], nums2: number[], k: number): number[] { + const n = nums1.length; + const arr = nums1.map((x, i) => [x, i]).sort((a, b) => a[0] - b[0]); + const pq = new MinPriorityQueue(); + let [s, j] = [0, 0]; + const ans: number[] = Array(k).fill(0); + for (let h = 0; h < n; ++h) { + const [x, i] = arr[h]; + while (j < h && arr[j][0] < x) { + const y = nums2[arr[j++][1]]; + pq.enqueue(y); + s += y; + if (pq.size() > k) { + s -= pq.dequeue(); + } + } + ans[i] = s; + } + return ans; +} ``` diff --git a/solution/3400-3499/3478.Choose K Elements With Maximum Sum/Solution.cpp b/solution/3400-3499/3478.Choose K Elements With Maximum Sum/Solution.cpp new file mode 100644 index 0000000000000..72f83b9eda258 --- /dev/null +++ b/solution/3400-3499/3478.Choose K Elements With Maximum Sum/Solution.cpp @@ -0,0 +1,30 @@ +class Solution { +public: + vector findMaxSum(vector& nums1, vector& nums2, int k) { + int n = nums1.size(); + vector> arr(n); + for (int i = 0; i < n; ++i) { + arr[i] = {nums1[i], i}; + } + ranges::sort(arr); + priority_queue, greater> pq; + long long s = 0; + int j = 0; + vector ans(n); + for (int h = 0; h < n; ++h) { + auto [x, i] = arr[h]; + while (j < h && arr[j].first < x) { + int y = nums2[arr[j].second]; + pq.push(y); + s += y; + if (pq.size() > k) { + s -= pq.top(); + pq.pop(); + } + ++j; + } + ans[i] = s; + } + return ans; + } +}; diff --git a/solution/3400-3499/3478.Choose K Elements With Maximum Sum/Solution.go b/solution/3400-3499/3478.Choose K Elements With Maximum Sum/Solution.go new file mode 100644 index 0000000000000..f607f0ba14d2c --- /dev/null +++ b/solution/3400-3499/3478.Choose K Elements With Maximum Sum/Solution.go @@ -0,0 +1,37 @@ +func findMaxSum(nums1 []int, nums2 []int, k int) []int64 { + n := len(nums1) + arr := make([][2]int, n) + for i, x := range nums1 { + arr[i] = [2]int{x, i} + } + ans := make([]int64, n) + sort.Slice(arr, func(i, j int) bool { return arr[i][0] < arr[j][0] }) + pq := hp{} + var s int64 + j := 0 + for h, e := range arr { + x, i := e[0], e[1] + for j < h && arr[j][0] < x { + y := nums2[arr[j][1]] + heap.Push(&pq, y) + s += int64(y) + if pq.Len() > k { + s -= int64(heap.Pop(&pq).(int)) + } + j++ + } + ans[i] = s + } + return ans +} + +type hp struct{ sort.IntSlice } + +func (h hp) Less(i, j int) bool { return h.IntSlice[i] < h.IntSlice[j] } +func (h *hp) Push(v any) { h.IntSlice = append(h.IntSlice, v.(int)) } +func (h *hp) Pop() any { + a := h.IntSlice + v := a[len(a)-1] + h.IntSlice = a[:len(a)-1] + return v +} diff --git a/solution/3400-3499/3478.Choose K Elements With Maximum Sum/Solution.java b/solution/3400-3499/3478.Choose K Elements With Maximum Sum/Solution.java new file mode 100644 index 0000000000000..5e9c00e3afcd5 --- /dev/null +++ b/solution/3400-3499/3478.Choose K Elements With Maximum Sum/Solution.java @@ -0,0 +1,28 @@ +class Solution { + public long[] findMaxSum(int[] nums1, int[] nums2, int k) { + int n = nums1.length; + int[][] arr = new int[n][0]; + for (int i = 0; i < n; ++i) { + arr[i] = new int[] {nums1[i], i}; + } + Arrays.sort(arr, (a, b) -> a[0] - b[0]); + PriorityQueue pq = new PriorityQueue<>(); + long s = 0; + long[] ans = new long[n]; + int j = 0; + for (int h = 0; h < n; ++h) { + int x = arr[h][0], i = arr[h][1]; + while (j < h && arr[j][0] < x) { + int y = nums2[arr[j][1]]; + pq.offer(y); + s += y; + if (pq.size() > k) { + s -= pq.poll(); + } + ++j; + } + ans[i] = s; + } + return ans; + } +} diff --git a/solution/3400-3499/3478.Choose K Elements With Maximum Sum/Solution.py b/solution/3400-3499/3478.Choose K Elements With Maximum Sum/Solution.py new file mode 100644 index 0000000000000..03a67f7ff6ff3 --- /dev/null +++ b/solution/3400-3499/3478.Choose K Elements With Maximum Sum/Solution.py @@ -0,0 +1,18 @@ +class Solution: + def findMaxSum(self, nums1: List[int], nums2: List[int], k: int) -> List[int]: + arr = [(x, i) for i, x in enumerate(nums1)] + arr.sort() + pq = [] + s = j = 0 + n = len(arr) + ans = [0] * n + for h, (x, i) in enumerate(arr): + while j < h and arr[j][0] < x: + y = nums2[arr[j][1]] + heappush(pq, y) + s += y + if len(pq) > k: + s -= heappop(pq) + j += 1 + ans[i] = s + return ans diff --git a/solution/3400-3499/3478.Choose K Elements With Maximum Sum/Solution.ts b/solution/3400-3499/3478.Choose K Elements With Maximum Sum/Solution.ts new file mode 100644 index 0000000000000..3b71a02edd5c5 --- /dev/null +++ b/solution/3400-3499/3478.Choose K Elements With Maximum Sum/Solution.ts @@ -0,0 +1,20 @@ +function findMaxSum(nums1: number[], nums2: number[], k: number): number[] { + const n = nums1.length; + const arr = nums1.map((x, i) => [x, i]).sort((a, b) => a[0] - b[0]); + const pq = new MinPriorityQueue(); + let [s, j] = [0, 0]; + const ans: number[] = Array(k).fill(0); + for (let h = 0; h < n; ++h) { + const [x, i] = arr[h]; + while (j < h && arr[j][0] < x) { + const y = nums2[arr[j++][1]]; + pq.enqueue(y); + s += y; + if (pq.size() > k) { + s -= pq.dequeue(); + } + } + ans[i] = s; + } + return ans; +}