3318. Find X-Sum of All K-Long Subarrays I
Description
You are given an array nums of n integers and two integers k and x.
The x-sum of an array is calculated by the following procedure:
- Count the occurrences of all elements in the array.
- Keep only the occurrences of the top
xmost frequent elements. If two elements have the same number of occurrences, the element with the bigger value is considered more frequent. - Calculate the sum of the resulting array.
Note that if an array has less than x distinct elements, its x-sum is the sum of the array.
Return an integer array answer of length n - k + 1 where answer[i] is the x-sum of the subarray nums[i..i + k - 1].
Example 1:
Input: nums = [1,1,2,2,3,4,2,3], k = 6, x = 2
Output: [6,10,12]
Explanation:
- For subarray
[1, 1, 2, 2, 3, 4], only elements 1 and 2 will be kept in the resulting array. Hence,answer[0] = 1 + 1 + 2 + 2. - For subarray
[1, 2, 2, 3, 4, 2], only elements 2 and 4 will be kept in the resulting array. Hence,answer[1] = 2 + 2 + 2 + 4. Note that 4 is kept in the array since it is bigger than 3 and 1 which occur the same number of times. - For subarray
[2, 2, 3, 4, 2, 3], only elements 2 and 3 are kept in the resulting array. Hence,answer[2] = 2 + 2 + 2 + 3 + 3.
Example 2:
Input: nums = [3,8,7,8,7,5], k = 2, x = 2
Output: [11,15,15,15,12]
Explanation:
Since k == x, answer[i] is equal to the sum of the subarray nums[i..i + k - 1].
Constraints:
1 <= n == nums.length <= 501 <= nums[i] <= 501 <= x <= k <= nums.length
Solutions
Solution 1: Hash Table + Ordered Set
We use a hash table $\textit{cnt}$ to count the occurrences of each element in the window, an ordered set $\textit{l}$ to store the $x$ elements with the highest occurrences in the window, and another ordered set $\textit{r}$ to store the remaining elements.
We maintain a variable $\textit{s}$ to represent the sum of the elements in $\textit{l}$. Initially, we add the first $k$ elements to the window, update the ordered sets $\textit{l}$ and $\textit{r}$, and calculate the value of $\textit{s}$. If the size of $\textit{l}$ is less than $x$ and $\textit{r}$ is not empty, we repeatedly move the largest element from $\textit{r}$ to $\textit{l}$ until the size of $\textit{l}$ equals $x$, updating the value of $\textit{s}$ in the process. If the size of $\textit{l}$ is greater than $x$, we repeatedly move the smallest element from $\textit{l}$ to $\textit{r}$ until the size of $\textit{l}$ equals $x$, updating the value of $\textit{s}$ in the process. At this point, we can calculate the current window's $\textit{x-sum}$ and add it to the answer array. Then we remove the left boundary element of the window, update $\textit{cnt}$, and update the ordered sets $\textit{l}$ and $\textit{r}$, as well as the value of $\textit{s}$. Continue traversing the array until the traversal is complete.
The time complexity is $O(n \times \log k)$, and the space complexity is $O(n)$. Here, $n$ is the length of the array $\textit{nums}$.
Similar problems:
Python3
class Solution:
def findXSum(self, nums: List[int], k: int, x: int) -> List[int]:
def add(v: int):
if cnt[v] == 0:
return
p = (cnt[v], v)
if l and p > l[0]:
nonlocal s
s += p[0] * p[1]
l.add(p)
else:
r.add(p)
def remove(v: int):
if cnt[v] == 0:
return
p = (cnt[v], v)
if p in l:
nonlocal s
s -= p[0] * p[1]
l.remove(p)
else:
r.remove(p)
l = SortedList()
r = SortedList()
cnt = Counter()
s = 0
n = len(nums)
ans = [0] * (n - k + 1)
for i, v in enumerate(nums):
remove(v)
cnt[v] += 1
add(v)
j = i - k + 1
if j < 0:
continue
while r and len(l) < x:
p = r.pop()
l.add(p)
s += p[0] * p[1]
while len(l) > x:
p = l.pop(0)
s -= p[0] * p[1]
r.add(p)
ans[j] = s
remove(nums[j])
cnt[nums[j]] -= 1
add(nums[j])
return ans
Java
class Solution {
private TreeSet<int[]> l = new TreeSet<>((a, b) -> a[0] == b[0] ? a[1] - b[1] : a[0] - b[0]);
private TreeSet<int[]> r = new TreeSet<>(l.comparator());
private Map<Integer, Integer> cnt = new HashMap<>();
private int s;
public int[] findXSum(int[] nums, int k, int x) {
int n = nums.length;
int[] ans = new int[n - k + 1];
for (int i = 0; i < n; ++i) {
int v = nums[i];
remove(v);
cnt.merge(v, 1, Integer::sum);
add(v);
int j = i - k + 1;
if (j < 0) {
continue;
}
while (!r.isEmpty() && l.size() < x) {
var p = r.pollLast();
s += p[0] * p[1];
l.add(p);
}
while (l.size() > x) {
var p = l.pollFirst();
s -= p[0] * p[1];
r.add(p);
}
ans[j] = s;
remove(nums[j]);
cnt.merge(nums[j], -1, Integer::sum);
add(nums[j]);
}
return ans;
}
private void remove(int v) {
if (!cnt.containsKey(v)) {
return;
}
var p = new int[] {cnt.get(v), v};
if (l.contains(p)) {
l.remove(p);
s -= p[0] * p[1];
} else {
r.remove(p);
}
}
private void add(int v) {
if (!cnt.containsKey(v)) {
return;
}
var p = new int[] {cnt.get(v), v};
if (!l.isEmpty() && l.comparator().compare(l.first(), p) < 0) {
l.add(p);
s += p[0] * p[1];
} else {
r.add(p);
}
}
}
C++
class Solution {
public:
vector<int> findXSum(vector<int>& nums, int k, int x) {
using pii = pair<int, int>;
set<pii> l, r;
int s = 0;
unordered_map<int, int> cnt;
auto add = [&](int v) {
if (cnt[v] == 0) {
return;
}
pii p = {cnt[v], v};
if (!l.empty() && p > *l.begin()) {
s += p.first * p.second;
l.insert(p);
} else {
r.insert(p);
}
};
auto remove = [&](int v) {
if (cnt[v] == 0) {
return;
}
pii p = {cnt[v], v};
auto it = l.find(p);
if (it != l.end()) {
s -= p.first * p.second;
l.erase(it);
} else {
r.erase(p);
}
};
vector<int> ans;
for (int i = 0; i < nums.size(); ++i) {
remove(nums[i]);
++cnt[nums[i]];
add(nums[i]);
int j = i - k + 1;
if (j < 0) {
continue;
}
while (!r.empty() && l.size() < x) {
pii p = *r.rbegin();
s += p.first * p.second;
r.erase(p);
l.insert(p);
}
while (l.size() > x) {
pii p = *l.begin();
s -= p.first * p.second;
l.erase(p);
r.insert(p);
}
ans.push_back(s);
remove(nums[j]);
--cnt[nums[j]];
add(nums[j]);
}
return ans;
}
};
Go