LeetCode-in-Java

3318. Find X-Sum of All K-Long Subarrays I

Easy

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:

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:

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:

Solution

import java.util.HashMap;
import java.util.Map;
import java.util.PriorityQueue;

public class Solution {
    private static class Pair {
        int num;
        int freq;

        Pair(int num, int freq) {
            this.num = num;
            this.freq = freq;
        }
    }

    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 - k + 1; i++) {
            HashMap<Integer, Integer> map = new HashMap<>();
            PriorityQueue<Pair> pq =
                    new PriorityQueue<>(
                            (a, b) -> {
                                if (a.freq == b.freq) {
                                    return b.num - a.num;
                                }
                                return b.freq - a.freq;
                            });
            for (int j = i; j < i + k; j++) {
                map.put(nums[j], map.getOrDefault(nums[j], 0) + 1);
            }
            for (Map.Entry<Integer, Integer> entry : map.entrySet()) {
                pq.add(new Pair(entry.getKey(), entry.getValue()));
            }
            int count = x;
            int sum = 0;
            while (!pq.isEmpty() && count > 0) {
                Pair pair = pq.remove();
                sum += pair.num * pair.freq;
                count--;
            }
            ans[i] = sum;
        }
        return ans;
    }
}