commit bbedc001ddf63b8c4db75767daa3fca558f46261
parent 4fd753a96fff65905a8159ae968acdaec1275654
Author: Jack Sangdahl <jack@pngu.org>
Date: Wed, 5 Nov 2025 09:47:25 -0700
3321. find x-sum of all k-long subarrays 2
Diffstat:
2 files changed, 469 insertions(+), 0 deletions(-)
diff --git a/find_x_sum_k_subarr2.c b/find_x_sum_k_subarr2.c
@@ -0,0 +1,377 @@
+// 3321. Find X-Sum of All K-Long Subarrays II
+#include "leetcode.h"
+
+/*
+ * you are given the array 'nums' of 'n' integers and two integers 'k' and 'x'.
+ * the x-sum of an array is calcluated by the following procedure: count the
+ * occurences of all elements in the array, keep only the occurences of the top
+ * 'x' most frequent elements. if two elements have the same number of
+ * occurences, the element with the bigger value is considered more frequent.
+ * calcluate the sum of the resulting array. note t hat 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]'
+ */
+
+struct item {
+ int id;
+ int freq;
+ long long val;
+};
+
+static inline int top_priority(const struct item *a, const struct item *b) {
+ if (a->freq != b->freq)
+ return a->freq < b->freq;
+ return a->val < b->val;
+}
+
+static inline int rest_priority(const struct item *a, const struct item *b) {
+ if (a->freq != b->freq)
+ return a->freq > b->freq;
+ return a->val > b->val;
+}
+
+struct heap {
+ struct item *a;
+ int size;
+ int capacity;
+ int (*priority)(const struct item *, const struct item *);
+};
+
+static void heap_init(struct heap *obj, int capacity,
+ int (*priority)(const struct item *,
+ const struct item *)) {
+ obj->capacity = capacity < 16 ? 16 : capacity;
+ obj->a = (struct item *)malloc(obj->capacity * sizeof(struct item));
+ obj->size = 0;
+ obj->priority = priority;
+}
+
+static void heap_free(struct heap *obj) {
+ free(obj->a);
+ obj->a = NULL;
+ obj->size = obj->capacity = 0;
+}
+
+static void heap_grow(struct heap *obj) {
+ if (obj->size + 4 >= obj->capacity) {
+ obj->capacity *= 2;
+ obj->a =
+ (struct item *)realloc(obj->a, obj->capacity * sizeof(struct item));
+ }
+}
+
+static void heap_push(struct heap *obj, struct item it) {
+ heap_grow(obj);
+ int i = obj->size++;
+ obj->a[i] = it;
+ while (i > 0) {
+ int p = (i - 1) >> 1;
+ if (obj->priority(&obj->a[i], &obj->a[p])) {
+ struct item tmp = obj->a[i];
+ obj->a[i] = obj->a[p];
+ obj->a[p] = tmp;
+ i = p;
+ } else
+ break;
+ }
+}
+
+static struct item heap_top(const struct heap *obj) { return obj->a[0]; }
+
+static int heap_empty(const struct heap *obj) { return !obj->size; }
+
+static void heap_pop(struct heap *obj) {
+ if (!obj->size)
+ return;
+ obj->a[0] = obj->a[--obj->size];
+ int i = 0;
+ while (true) {
+ int l = i * 2 + 1, r = l + 1, best = i;
+ if (l < obj->size && obj->priority(&obj->a[l], &obj->a[best]))
+ best = l;
+ if (r < obj->size && obj->priority(&obj->a[r], &obj->a[best]))
+ best = r;
+ if (best == i)
+ break;
+ struct item tmp = obj->a[i];
+ obj->a[i] = obj->a[best];
+ obj->a[best] = tmp;
+ i = best;
+ }
+}
+
+static int cmp(const void *a, const void *b) {
+ long long aa = *(const long long *)a, bb = *(const long long *)b;
+ if (aa < bb)
+ return -1;
+ if (aa > bb)
+ return 1;
+ return 0;
+}
+
+static int val_to_id(long long *unique, int m, long long val) {
+ int l = 0, r = m - 1;
+ while (l <= r) {
+ int mid = (l + r) >> 1;
+ if (unique[mid] == val)
+ return mid;
+ if (unique[mid] < val)
+ l = mid + 1;
+ else
+ r = mid - 1;
+ }
+ return -1;
+}
+
+long long *findXSum(int *nums, int numsSize, int k, int x, int *returnSize) {
+ if (!numsSize) {
+ *returnSize = 0;
+ return NULL;
+ }
+ *returnSize = numsSize - k + 1;
+ long long *ans = (long long *)malloc(*returnSize * sizeof(long long));
+ if (k == x) {
+ long long s = 0;
+ for (int i = 0; i < k; ++i)
+ s += nums[i];
+ ans[0] = s;
+ for (int i = k; i < numsSize; ++i) {
+ s += nums[i] - nums[i - k];
+ ans[i - k + 1] = s;
+ }
+ return ans;
+ }
+ long long *tmp = (long long *)malloc(numsSize * sizeof(long long));
+ for (int i = 0; i < numsSize; ++i)
+ tmp[i] = (long long)nums[i];
+ qsort(tmp, numsSize, sizeof(long long), cmp);
+ int m = 0;
+ for (int i = 0; i < numsSize; ++i)
+ if (!i || tmp[i] != tmp[i - 1])
+ tmp[m++] = tmp[i];
+ long long *unique = tmp;
+ int *freq = (int *)calloc(m, sizeof(int));
+ char *in_top = (char *)calloc(m, sizeof(char));
+ struct heap top, rest;
+ heap_init(&top, m * 2 + 16, top_priority);
+ heap_init(&rest, m * 2 + 16, rest_priority);
+ for (int i = 0; i < k; ++i) {
+ int id = val_to_id(unique, m, nums[i]);
+ freq[id]++;
+ }
+ for (int i = 0; i < m; ++i)
+ if (freq[i] > 0) {
+ struct item it = {i, freq[i], unique[i]};
+ heap_push(&rest, it);
+ }
+ long long x_sum = 0;
+ int in_top_cnt = 0;
+ while (in_top_cnt < x) {
+ while (!heap_empty(&rest)) {
+ struct item it = heap_top(&rest);
+ int id = it.id;
+ if (it.freq != freq[id] || it.freq <= 0 || in_top[id]) {
+ heap_pop(&rest);
+ continue;
+ }
+ heap_pop(&rest);
+ in_top[id] = 1;
+ ++in_top_cnt;
+ x_sum += (long long)freq[id] * unique[id];
+ struct item new_it = {id, freq[id], unique[id]};
+ heap_push(&top, new_it);
+ break;
+ }
+ if (heap_empty(&rest))
+ break;
+ }
+ while (true) {
+ while (!heap_empty(&top)) {
+ struct item it = heap_top(&top);
+ int id = it.id;
+ if (it.freq != freq[id] || it.freq <= 0 || !in_top[id]) {
+ heap_pop(&top);
+ continue;
+ }
+ break;
+ }
+ while (!heap_empty(&rest)) {
+ struct item it = heap_top(&rest);
+ int id = it.id;
+ if (it.freq != freq[id] || it.freq <= 0 || in_top[id]) {
+ heap_pop(&rest);
+ continue;
+ }
+ break;
+ }
+ if (heap_empty(&top) || heap_empty(&rest))
+ break;
+ struct item worst = heap_top(&top);
+ struct item best = heap_top(&rest);
+ int worst_id = worst.id, best_id = best.id;
+ if (!(worst.freq == freq[worst_id] && in_top[worst_id] &&
+ best.freq == freq[best_id] && !in_top[best_id])) {
+ if (!(worst.freq == freq[worst_id] && in_top[worst_id]))
+ heap_pop(&top);
+ if (!(best.freq == freq[best_id] && !in_top[best_id]))
+ heap_pop(&rest);
+ continue;
+ }
+ int do_swap = 0;
+ if (best.freq > worst.freq)
+ do_swap = 1;
+ else if (best.freq == worst.freq && unique[best_id] > unique[worst_id])
+ do_swap = 0;
+ if (!do_swap)
+ break;
+ heap_pop(&top);
+ heap_pop(&rest);
+ in_top[worst_id] = 0;
+ in_top[best_id] = 1;
+ x_sum -= (long long)worst.freq * unique[worst_id];
+ x_sum += (long long)best.freq * unique[best_id];
+ struct item new_worst = {worst_id, freq[worst_id], unique[worst_id]};
+ struct item new_best = {best_id, freq[best_id], unique[best_id]};
+ heap_push(&rest, new_worst);
+ heap_push(&top, new_best);
+ }
+ ans[0] = x_sum;
+ for (int i = k; i < numsSize; ++i) {
+ long long out_val = nums[i - k];
+ long long in_val = nums[i];
+ int out_id = val_to_id(unique, m, out_val);
+ int in_id = val_to_id(unique, m, in_val);
+ if (in_top[out_id])
+ x_sum -= unique[out_id];
+ freq[out_id]--;
+ if (freq[out_id] > 0) {
+ struct item new_it = {out_id, freq[out_id], unique[out_id]};
+ heap_push(&rest, new_it);
+ heap_push(&top, new_it);
+ } else {
+ if (in_top[out_id]) {
+ in_top[out_id] = 0;
+ --in_top_cnt;
+ }
+ }
+ freq[in_id]++;
+ if (in_top[in_id])
+ x_sum += unique[in_id];
+ struct item in_it = {in_id, freq[in_id], unique[in_id]};
+ heap_push(&rest, in_it);
+ heap_push(&top, in_it);
+ while (in_top_cnt < x) {
+ while (!heap_empty(&rest)) {
+ struct item it = heap_top(&rest);
+ int id = it.id;
+ if (it.freq != freq[id] || it.freq <= 0 || in_top[id]) {
+ heap_pop(&rest);
+ continue;
+ }
+ heap_pop(&rest);
+ in_top[id] = 1;
+ ++in_top_cnt;
+ x_sum += (long long)freq[id] * unique[id];
+ struct item new_it = {id, freq[id], unique[id]};
+ heap_push(&top, new_it);
+ break;
+ }
+ if (heap_empty(&rest))
+ break;
+ }
+ while (true) {
+ while (!heap_empty(&top)) {
+ struct item it = heap_top(&top);
+ int id = it.id;
+ if (it.freq != freq[id] || it.freq <= 0 || !in_top[id]) {
+ heap_pop(&top);
+ continue;
+ }
+ break;
+ }
+ while (!heap_empty(&rest)) {
+ struct item it = heap_top(&rest);
+ int id = it.id;
+ if (it.freq != freq[id] || it.freq <= 0 || in_top[id]) {
+ heap_pop(&rest);
+ continue;
+ }
+ break;
+ }
+ if (heap_empty(&top) || heap_empty(&rest))
+ break;
+ struct item worst = heap_top(&top);
+ struct item best = heap_top(&rest);
+ int worst_id = worst.id, best_id = best.id;
+ if (!(worst.freq == freq[worst_id] && in_top[worst_id] &&
+ best.freq == freq[best_id] && !in_top[best_id])) {
+ if (!(worst.freq == freq[worst_id] && in_top[worst_id]))
+ heap_pop(&top);
+ if (!(best.freq == freq[best_id] && !in_top[best_id]))
+ heap_pop(&rest);
+ continue;
+ }
+ int do_swap = 0;
+ if (best.freq > worst.freq)
+ do_swap = 1;
+ else if (best.freq == worst.freq && unique[best_id] > unique[worst_id])
+ do_swap = 1;
+ if (!do_swap)
+ break;
+ heap_pop(&top);
+ heap_pop(&rest);
+ in_top[worst_id] = 0;
+ in_top[best_id] = 1;
+ x_sum -= (long long)worst.freq * unique[worst_id];
+ x_sum += (long long)best.freq * unique[best_id];
+ struct item new_worst = {worst_id, freq[worst_id], unique[worst_id]};
+ struct item new_best = {best_id, freq[best_id], unique[best_id]};
+ heap_push(&rest, new_worst);
+ heap_push(&top, new_best);
+ }
+ if (in_top_cnt > x) {
+ while (in_top_cnt > x) {
+ while (!heap_empty(&top)) {
+ struct item it = heap_top(&top);
+ int id = it.id;
+ if (it.freq != freq[id] || it.freq <= 0 || !in_top[id]) {
+ heap_pop(&top);
+ continue;
+ }
+ heap_pop(&top);
+ in_top[id] = 0;
+ --in_top_cnt;
+ x_sum -= (long long)it.freq * unique[id];
+ struct item new_it = {id, freq[id], unique[id]};
+ heap_push(&rest, new_it);
+ break;
+ }
+ if (heap_empty(&top))
+ break;
+ }
+ }
+ ans[i - k + 1] = x_sum;
+ }
+ free(freq);
+ free(in_top);
+ free(unique);
+ heap_free(&top);
+ heap_free(&rest);
+ return ans;
+}
+
+int main() {
+ int n1[] = {1, 1, 2, 2, 3, 4, 2, 3}, rs1;
+ int n2[] = {3, 8, 7, 8, 7, 5}, rs2;
+ long long *fxs1 = findXSum(n1, ARRAY_SIZE(n1), 6, 2, &rs1);
+ long long *fxs2 = findXSum(n2, ARRAY_SIZE(n2), 2, 2, &rs2);
+ for (int i = 0; i < rs1; i++)
+ printf("%lld ", fxs1[i]); // expect: 6 10 12
+ printf("\n");
+ for (int i = 0; i < rs2; i++)
+ printf("%lld ", fxs2[i]); // expect: 11 15 15 15 12
+ printf("\n");
+ free(fxs1);
+ free(fxs2);
+}
diff --git a/find_x_sum_k_subarr2.py b/find_x_sum_k_subarr2.py
@@ -0,0 +1,92 @@
+# 3321. Find X-Sum of All K-Long Subarrays II
+from collections import Counter
+
+from sortedcontainers import SortedList
+
+"""
+you are given the array 'nums' of 'n' integers and two integers 'k' and 'x'.
+the x-sum of an array is calcluated by the following procedure: count the
+occurences of all elements in the array, keep only the occurences of the top
+'x' most frequent elements. if two elements have the same number of
+occurences, the element with the bigger value is considered more frequent.
+calcluate the sum of the resulting array. note t hat 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]'
+"""
+
+
+class Solution(object):
+ def findXSum(self, nums, k, x):
+ """
+ :type nums: List[int]
+ :type k: int
+ :type x: int
+ :rtype: List[int]
+ """
+ freq = Counter()
+ top = SortedList()
+ rest = SortedList()
+ top_sum, ans = 0, []
+
+ def balance():
+ nonlocal top_sum
+ while len(top) < x and rest:
+ f, v = rest.pop()
+ top.add((f, v))
+ top_sum += f * v
+ while len(top) > x:
+ f, v = top.pop(0)
+ top_sum -= f * v
+ rest.add((f, v))
+ while rest and top and rest[-1] > top[0]:
+ f1, v1 = rest.pop()
+ f2, v2 = top.pop(0)
+ top_sum += f1 * v1 - f2 * v2
+ top.add((f1, v1))
+ rest.add((f2, v2))
+
+ def add(num):
+ nonlocal top_sum
+ old = (freq[num], num)
+ if old in top:
+ top.remove(old)
+ top_sum -= old[0] * old[1]
+ elif old in rest:
+ rest.remove(old)
+ freq[num] += 1
+ new = (freq[num], num)
+ rest.add(new)
+ balance()
+
+ def remove(num):
+ nonlocal top_sum
+ old = (freq[num], num)
+ if old in top:
+ top.remove(old)
+ top_sum -= old[0] * old[1]
+ else:
+ rest.remove(old)
+ freq[num] -= 1
+ if freq[num] > 0:
+ rest.add((freq[num], num))
+ else:
+ del freq[num]
+ balance()
+
+ for i in range(k):
+ add(nums[i])
+ ans.append(top_sum)
+
+ for i in range(k, len(nums)):
+ remove(nums[i - k])
+ add(nums[i])
+ ans.append(top_sum)
+
+ return ans
+
+
+if __name__ == "__main__":
+ obj = Solution()
+ print(obj.findXSum(nums=[1, 1, 2, 2, 3, 4, 2, 3], k=6, x=2))
+ print(obj.findXSum(nums=[3, 8, 7, 8, 7, 5], k=2, x=2))
