3351. Good Subsequences Sum

Sum of Good Subsequences

Counting and summing subsequences share a lot in spirit with doing the same for subarrays.

Before diving into problem 3351, let me briefly cover similarities and differences.

Subarrays vs. Subsequences: Similar but Different

A. Similar: Each Visited Element Tries to Be the "Tail"

When searching for subarrays, we let each visited index $i$ try to serve as subarray's tail,

asking how many subarrays end at index $i$.

By the same logic, we can let each visited index $i$ serve as subsequence's tail.

This keeps the loop logic clean and avoids double-counting.

B. Different: How to Extend and Pass Down the Lineage

One of the main divergences. Subarrays require strict index continuity:

index $i$ can only extend subarrays that end at index $i - 1$.

But subsequences allow skipping indices. Any relative order that's preserved is valid:

index $i$ can connect to any subsequence ending at indices $0, \ldots, i - 1$.

The Pattern in Problem 3351

For each currently visited index $i$, we look at how many subsequences with tail value $nums[i] - 1$ or $nums[i] + 1$ have appeared in the previous $i$ iterations,

since the problem states that adjacent elements in each subsequence must have an absolute difference of exactly 1.

Say among the previous $i$ iterations there are:

I. $j$ subsequences with tail value $nums[i] - 1$, with total sum $Sum_{nums[i] - 1}$

II. $k$ subsequences with tail value $nums[i] + 1$, with total sum $Sum_{nums[i] + 1}$

We can immediately conclude two facts:

1. The number of subsequences with tail at index $i$ is $j + k + 1$

2. The number of subsequences with tail value of $nums[i]$ net increases by $j + k + 1$

Note that statements 1 and 2 use slightly different phrasing, because there is only one index $i$,

but before index $i$, subsequences with tail value of $nums[i]$ may have already formed.

Take a moment to digest the difference between these two statements before you read on~~

-------------No rush, take your time-------------

Back to $j + k + 1$. Where does $+ 1$ come from?

Don't forget: at index $i$, we also have the right to start fresh, letting $nums[i]$ form a subsequence on its own.

With counting handled, summing follows naturally. We now have:

(1). Count of subsequences with tail value of $nums[i]$ net increases by $j + k + 1$

(2). Total sum of subsequences with tail value of $nums[i]$ net increases by

$(Sum_{nums[i] - 1} + nums[i] \times j) + (Sum_{nums[i] + 1} + nums[i] \times k) + (nums[i])$

This is a bit tricky. I've used three pairs of parentheses to explain each part 😄

a. First bracket: index $i$ extends all subsequences with tail value of $nums[i] - 1$.

Index $i$ inherits their total sum $Sum_{nums[i] - 1}$,

and $nums[i]$ is appended to all $j$ subsequences with tail value of $nums[i] - 1$.

b. Middle bracket: index $i$ extends all subsequences with tail value of $nums[i] + 1$.

Index $i$ inherits their total sum $Sum_{nums[i] + 1}$,

and $nums[i]$ is appended to all $k$ subsequences with tail value of $nums[i] + 1$.

c. Last bracket: index $i$ lets $nums[i]$ form a standalone subsequence.

So parts a, b, and c together give the net increase in total sum for subsequences with tail value of $nums[i]$.

After the full traversal, we sum up the total sums across all distinct tail values,

then apply modulo $10^9 + 7$ as required.

One thing to note: the number of subsequences is on the order of $O(2^n)$,

and throughout our pipeline, we use subsequence counts to compute subsequence sums,

so we must apply modulo to both counts and sums to prevent overflow.

Time and space complexity: both $O(n)$.

C++

#include <unordered_map>
#include <vector>
using namespace std;

int sumGoodSubsequences(vector<int>& nums) // LeetCode Q.3351.
{
    long long modulo = pow(10, 9) + 7; // Long long prevents overflow.

    unordered_map<int, long long> subseqCounts, subseqSums;

    for (const auto& num : nums) {
        for (const auto& validTail : {num - 1, num + 1}) {
            long long prevCount = subseqCounts[validTail];
            subseqCounts[num] += prevCount;

            long long prevSum = subseqSums[validTail];
            subseqSums[num] += prevSum + num * prevCount;
        }

        // Counts & sums are O(2^n) so must do modulo for future operations.
        subseqCounts[num] += 1;
        if (subseqCounts[num] > modulo)
            subseqCounts[num] %= modulo;

        subseqSums[num] += num;
        if (subseqSums[num] > modulo)
            subseqSums[num] %= modulo;
    }

    long long goodSubseqSum = 0;
    for (const auto& pair : subseqSums)
        goodSubseqSum += pair.second;

    return goodSubseqSum % modulo; // Control value size.
}

Python

def sum_good_subsequences(nums: list[int]) -> int:  # LeetCode Q.3351.
    modulo = 10 ** 9 + 7  # Prevents overflow.

    unique_nums = set(nums)

    subseqCounts: dict[int, int] = dict(
        zip(unique_nums, [0] * len(unique_nums))
    )

    subseqSums: dict[int, int] = dict(
        zip(unique_nums, [0] * len(unique_nums))
    )

    for num in nums:
        for validTail in (num - 1, num + 1):
            prev_count = subseqCounts.get(validTail, 0)
            subseqCounts[num] += prev_count

            prev_sum = subseqSums.get(validTail, 0)
            subseqSums[num] += prev_sum + num * prev_count

        # Counts & sums are O(2^n) so must do modulo for future operations.
        subseqCounts[num] += 1
        if subseqCounts[num] > modulo: subseqCounts[num] %= modulo

        subseqSums[num] += num
        if subseqSums[num] > modulo: subseqSums[num] %= modulo

    return sum(subseqSums.values()) % modulo  # Control value magnitude.

Follow-up Problem

If the problem asked for subarrays instead of subsequences, how would the optimal time and space complexity change?