Subsequence sum after capping elements

Array Two Pointers Dynamic Programming Sorting

25.1% acceptance

Feb 25, 2026

161

You are given an integer array nums of size n and a positive integer k. An array capped by value x is obtained by replacing every element nums[i] with min(nums[i], x). For each integer x from 1 to n, determine whether it is possible to choose a subsequence from the array capped by x such that the sum of the chosen elements is exactly k. Return a 0-indexed boolean array answer of size n, where answer[i] is true if it is possible when using x = i + 1, and false otherwise.

Solution

Rust

Time O(n³)

Space O(n)

LeetCode

solution.rs

impl Solution {
  pub fn subsequence_sum_after_capping(nums: Vec<i32>, k: i32) -> Vec<bool> {
    let n = nums.len();
    let k = k as usize;
    // Key insight (O(n*k) total):
    //   For cap x, the capped array = {true values for nums[i]<=x} union {x repeated C(x) times},
    //   where C(x) = #{i : nums[i] > x}.
    //   All "capped" items share the same value x, so their contribution counts are
    //   0*x, 1*x, ..., C(x)*x -- a simple arithmetic sequence, no need for a 2-D DP.
    //
    //   Algorithm:
    //     1. Build dp_fixed incrementally: dp_fixed[s] = can we pick a subset of
    //        {i : nums[i] <= x} (at true values) summing to s?
    //        Elements unlock in value order: add freq[x] copies at each step -> O(n*k) total.
    //     2. answer[x-1] = exists j in {0..=C(x)} s.t. k-j*x >= 0 && dp_fixed[k-j*x].
    //        Loop runs O(k/x) per x, total O(k log n).
    //   Total: O(n*k).

    // freq[v] = #{i : nums[i] == v} (values > n always stay capped, handled by cap_count)
    let mut freq = vec![0usize; n + 2];
    for &v in &nums {
      let v = v as usize;
      if v <= n {
        freq[v] += 1;
      }
    }

    // prefix[v] = #{i : nums[i] <= v}
    let mut prefix = vec![0usize; n + 2];
    for v in 1..=n {
      prefix[v] = prefix[v - 1] + freq[v];
    }

    // dp[s]: can we pick a subset of {i : nums[i] <= x} (true values) summing to s?
    let mut dp = vec![false; k + 1];
    dp[0] = true;

    let mut answer = vec![false; n];

    for x in 1..=n {
      // Add freq[x] copies of value x (0/1 knapsack, one copy at a time)
      for _ in 0..freq[x] {
        for s in (x..=k).rev() {
          if dp[s - x] {
            dp[s] = true;
          }
        }
      }

      // cap_count(x) = #{i : nums[i] > x}; each contributes value x when capped
      let cap_count = n - prefix[x];

      // Check: exists j in {0..=min(cap_count, k/x)} s.t. dp[k - j*x]
      let max_j = cap_count.min(k / x);
      for j in 0..=max_j {
        if dp[k - j * x] {
          answer[x - 1] = true;
          break;
        }
      }
    }

    answer
  }
}