Find x sum of all k long subarrays ii

Array Hash Table Sliding Window Heap (Priority Queue)

41.1% acceptance

Feb 23, 2026

448

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 x most 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].

Solution

Rust

Time O(n log n)

Space O(n)

LeetCode

solution.rs

use std::collections::{BTreeSet, HashMap};

impl Solution {
  pub fn find_x_sum(nums: Vec<i32>, k: i32, x: i32) -> Vec<i64> {
    let k = k as usize;
    let x = x as usize;
    let n = nums.len();

    let mut freq: HashMap<i32, i64> = HashMap::new();
    // BTreeSet<(count, val)>: max element = most frequent (tie-break: larger val)
    let mut top: BTreeSet<(i64, i32)> = BTreeSet::new(); // top-x distinct elements
    let mut bot: BTreeSet<(i64, i32)> = BTreeSet::new(); // remaining distinct elements
    let mut top_sum: i64 = 0;
    if n < k { return vec![]; }
    let mut result = Vec::with_capacity(n - k + 1);

    for i in 0..n {
      // Add nums[i] to the window
      let v = nums[i];
      let old_cnt = *freq.get(&v).unwrap_or(&0);
      let new_cnt = old_cnt + 1;
      freq.insert(v, new_cnt);

      if old_cnt > 0 {
        if top.remove(&(old_cnt, v)) {
          top_sum -= old_cnt * v as i64;
          top.insert((new_cnt, v));
          top_sum += new_cnt * v as i64;
        } else {
          bot.remove(&(old_cnt, v));
          bot.insert((new_cnt, v));
        }
      } else {
        bot.insert((1, v));
      }

      Self::rebalance(&mut top, &mut bot, &mut top_sum, x);

      if i >= k {
        // Remove nums[i - k] from the window
        let v = nums[i - k];
        let old_cnt = *freq.get(&v).unwrap();
        let new_cnt = old_cnt - 1;
        if new_cnt == 0 {
          freq.remove(&v);
        } else {
          freq.insert(v, new_cnt);
        }

        if top.remove(&(old_cnt, v)) {
          top_sum -= old_cnt * v as i64;
          if new_cnt > 0 {
            bot.insert((new_cnt, v));
          }
        } else {
          bot.remove(&(old_cnt, v));
          if new_cnt > 0 {
            bot.insert((new_cnt, v));
          }
        }

        Self::rebalance(&mut top, &mut bot, &mut top_sum, x);
      }

      if i >= k - 1 {
        result.push(top_sum);
      }
    }

    result
  }

  fn rebalance(
    top: &mut BTreeSet<(i64, i32)>,
    bot: &mut BTreeSet<(i64, i32)>,
    top_sum: &mut i64,
    x: usize,
  ) {
    // Move excess from top to bot
    while top.len() > x {
      let e = *top.iter().next().unwrap();
      top.remove(&e);
      *top_sum -= e.0 * e.1 as i64;
      bot.insert(e);
    }
    // Fill top from bot if underfull
    while top.len() < x && !bot.is_empty() {
      let e = *bot.iter().next_back().unwrap();
      bot.remove(&e);
      top.insert(e);
      *top_sum += e.0 * e.1 as i64;
    }
    // Single swap: ensure all top >= all bot (order invariant)
    if !top.is_empty() && !bot.is_empty() {
      let t_min = *top.iter().next().unwrap();
      let b_max = *bot.iter().next_back().unwrap();
      if b_max > t_min {
        top.remove(&t_min);
        *top_sum -= t_min.0 * t_min.1 as i64;
        bot.remove(&b_max);
        top.insert(b_max);
        *top_sum += b_max.0 * b_max.1 as i64;
        bot.insert(t_min);
      }
    }
  }
}