Lab 12 — SOS DP (Sum over Subsets)

Goal

Compute F[mask] = Σ A[sub] over all sub ⊆ mask, for all masks simultaneously, in O(2ⁿ · n) — replacing the O(4ⁿ) all-pairs-of-masks or O(3ⁿ) subset-enumeration approaches. Apply it to LeetCode 982 (Triples with Bitwise AND Equal To Zero). After this lab you should write the forward transform, its Möbius inverse, and the 3ⁿ subset loop from memory, and recognize when each wins.

Background Concepts

The naive “for each mask, sum over its subsets” costs O(3ⁿ) (each of 2ⁿ masks enumerates its subsets; Σ over masks of 2^popcount = 3ⁿ). SOS DP does it in O(2ⁿ · n) by a dimension-by-dimension relaxation — the subset-sum (zeta) transform:

def sos(A, n):
    F = A[:]
    for i in range(n):                     # one bit-dimension at a time
        bit = 1 << i
        for mask in range(1 << n):
            if mask & bit:
                F[mask] += F[mask ^ bit]
    return F

Interpretation as n-dimensional prefix sum. Think of a mask as a point in the hypercube {0,1}ⁿ. Dimension i sweeps the “prefix sum along axis i”: after processing bit i, F[mask] holds the sum over all sub-masks that agree with mask on bits > i (already processed) and are ≤ mask on bits ≤ i. After all n sweeps, F[mask] = sum over every sub ⊆ mask. It is exactly the standard multi-dimensional prefix-sum, one axis per bit, each axis of length 2.

The 3ⁿ subset-enumeration trick is the intermediate power tool:

sub = mask
while True:
    # visit sub ⊆ mask
    if sub == 0: break
    sub = (sub - 1) & mask                  # next lower submask

Use 3ⁿ when you need only one mask’s subset sum, or when the per-subset work is non-additive. Use SOS’s 2ⁿ·n when you need the sums for all masks. Crossover: for n ≈ 20, 3ⁿ ≈ 3.5×10⁹ vs 2ⁿ·n ≈ 2×10⁷ — SOS wins by two orders of magnitude when all masks are needed.

Inverse (Möbius) transform. Undo the zeta by subtracting in the same sweep structure: F[mask] -= F[mask ^ bit]. Zeta ∘ Möbius = identity. This is the inclusion-exclusion inverse.

Superset sums are the mirror image: process the unset bit direction (F[mask] += F[mask | bit]), or equivalently run subset SOS on complemented indices.

Interview Context

SOS DP is niche but a hard differentiator: it appears in bitmask-heavy counting problems where a candidate’s first instinct is O(4ⁿ) pair enumeration. Recognizing “I need, for every mask, the aggregate over all its submasks” and reaching for the 2ⁿ·n transform is a senior/competitive signal. LeetCode has few pure hits — the honest framing is that SOS is a technique you fold into subset-DP problems (Maximum Students Taking Exam, Count Subtrees) and AND/OR pair-counting. LC 982 is the cleanest canonical showcase.

Problem Statement (LC 982)

Given an integer array nums, return the number of triples of indices (i, j, k) (each independently ranging over all indices, repetition allowed, order counted) such that nums[i] & nums[j] & nums[k] == 0, where & is bitwise AND.

Constraints

  • 1 ≤ nums.length ≤ 1000.
  • 0 ≤ nums[i] < 2¹⁶ (so masks live in a 16-bit universe, 2¹⁶ = 65536).

Clarifying Questions

  1. Are i, j, k ordered and independently chosen? (Yes — (0,1,2) and (2,1,0) are distinct; i = j is allowed.)
  2. What’s the bit width? (Values < 2¹⁶, so a 16-bit mask universe — U = 2¹⁶.)
  3. Do we count value-triples or index-triples? (Index-triples — duplicates in nums each count.)
  4. Is the answer possibly large? (Up to 1000³ = 10⁹ — fits in 64-bit, but note it in languages with 32-bit ints.)
  5. Is AND == 0 a subset or a disjointness condition? (x & y & z == 0 means the ANDed bits are a subset of the complement of… — see the reduction below.)

Examples

nums = [2, 1, 3] → 12
nums = [0, 0, 0] → 27   (every triple ANDs to 0)

Initial Brute Force

Triple nested loop over all (i, j, k), test nums[i] & nums[j] & nums[k] == 0, count. O(m³) with m = len(nums) — the oracle.

Brute Force Complexity

O(m³). At m = 1000 that is 10⁹ AND-and-compare operations — borderline-to-TLE in Python, and the pattern generalizes badly. A better structure: precompute a histogram of pairwise ANDs (pairs[v] = number of (i, j) with nums[i] & nums[j] == v) in O(m²), then for each k count pairs whose AND is disjoint from nums[k].

Optimization Path

The reduction: x & z == 0 ⟺ x ⊆ ~z (every set bit of x is unset in z). So nums[i] & nums[j] & nums[k] == 0 ⟺ (nums[i] & nums[j]) ⊆ ~nums[k]. Therefore:

  1. Build pairs[v] = #{(i, j) : nums[i] & nums[j] == v} in O(m²).
  2. SOS: F[mask] = Σ_{v ⊆ mask} pairs[v] = number of pairs whose AND is a submask of mask, for all masks in O(U · log U).
  3. Answer = Σ_k F[~nums[k] & (U−1)] — for each k, count pairs disjoint from nums[k].

Total O(m² + U log U) instead of O(m³). The SOS step is what turns “for each k, re-scan all pair-values” (another O(U) or O(m²)) into an O(1) table lookup per k.

Final Expected Approach

def sos(A, n):
    """F[mask] = Σ A[sub] over all sub ⊆ mask, all masks, in O(2ⁿ · n)."""
    F = A[:]
    for i in range(n):
        bit = 1 << i
        for mask in range(1 << n):
            if mask & bit:                 # relax along bit-dimension i
                F[mask] += F[mask ^ bit]
    return F

def inverse_sos(F, n):
    """Möbius inverse: recover A from F = sos(A). Subtract in the same order."""
    A = F[:]
    for i in range(n):
        bit = 1 << i
        for mask in range(1 << n):
            if mask & bit:
                A[mask] -= A[mask ^ bit]
    return A

def count_and_triples(nums):
    """LC 982: # ordered triples (i,j,k) with nums[i] & nums[j] & nums[k] == 0."""
    B = 16
    FULL = (1 << B) - 1
    pairs = [0] * (1 << B)
    for x in nums:                         # O(m²) pairwise-AND histogram
        for y in nums:
            pairs[x & y] += 1
    F = sos(pairs, B)                      # F[m] = #pairs with AND ⊆ m
    return sum(F[FULL ^ z] for z in nums)  # AND & z == 0  ⟺  AND ⊆ ~z

# 3ⁿ single-mask subset sum (the intermediate power tool)
def subset_sum_of(A, mask):
    total, sub = 0, mask
    while True:
        total += A[sub]
        if sub == 0:
            break
        sub = (sub - 1) & mask             # next lower submask of mask
    return total

Data Structures Used

  • A flat array F of length 2ⁿ indexed by bitmask — the transform is entirely in-place-ish over it.
  • The pairs[] histogram of size 2¹⁶ for LC 982.
  • Bit operations: 1 << i, mask ^ bit, (sub − 1) & mask, ~z & FULL for the complement.

Correctness Argument

Loop invariant of the zeta transform. After processing dimensions 0..i−1, F[mask] equals Σ A[sub] over all sub that equal mask on bits ≥ i and are arbitrary (≤ mask) on bits < i. Base: before any dimension, F[mask] = A[mask] (only sub = mask). Step: when we process bit i for a mask with bit i set, F[mask] += F[mask ^ bit] fuses the “bit-i-set” family with the “bit-i-clear” family, both already summed over bits < i — extending the invariant to include bit i. After all n dimensions, the free set is all bits, i.e. every sub ⊆ mask. Each sub is added exactly once because each dimension merges disjoint families.

Möbius inverts zeta. Subtracting A[mask ^ bit] in the same sweep order peels off precisely the contribution the forward pass added, so inverse_sos(sos(A)) == A — the standard subset-lattice inclusion-exclusion pair.

LC 982 reduction. a & z == 0 ⟺ a ⊆ ~z. F[~z] sums pairs[a] over all a ⊆ ~z, i.e. counts pairs (i, j) whose AND is disjoint from z = nums[k]. Summing over k counts all valid ordered triples.

Complexity

StepTimeSpace
Pair histogram (LC 982)O(m²)O(2ⁿ)
SOS forward / inverseO(2ⁿ · n)O(2ⁿ)
Answer aggregationO(m)O(1)
3ⁿ single-mask subset sumO(3ⁿ) total over all masksO(1)

Implementation Requirements

  • Copy the input (F = A[:]) — the transform mutates in place; don’t clobber the caller’s array unless intended.
  • Iterate dimensions outer, masks inner; swapping the loop order breaks the prefix-sum semantics.
  • Only update F[mask] when mask & bit — skipping this doubles-adds.
  • For LC 982, use the full 16-bit universe even if max(nums) is small; sizing to max+1 bits is fine but must cover every complement ~z.

Tests

  • sos(A, n) vs the 3ⁿ subset-enumeration sum for every mask, random A, n ∈ {3,4,5}.
  • inverse_sos(sos(A, n), n) == A for random A.
  • LC 982: [2,1,3] → 12, [0,0,0] → 27.
  • LC 982 vs O(m³) brute force on random arrays (values up to 2¹⁶ and small values), dozens of trials.
  • Superset-sum variant: verify F[mask] += F[mask | bit] gives Σ over supersets.

Follow-up Questions

  • What are the zeta and Möbius transforms, and how do they connect to inclusion-exclusion? The subset-sum transform is the zeta transform on the subset lattice; its inverse is the Möbius transform, whose kernel is (−1)^popcount(mask \ sub). Alternating-sign subset sums are inclusion-exclusion — the same machinery as Phase 12’s inclusion-exclusion lab, specialized to the boolean lattice where Möbius is just “subtract each dropped bit.”
  • How do you convolve over subsets — C[mask] = Σ_{s∪t=mask, s∩t=∅} A[s]·B[t]? Subset-sum convolution: rank the arrays by popcount into n+1 layers, zeta-transform each layer, multiply layer-wise with a popcount-carrying convolution, then Möbius back — O(2ⁿ · n²). It fixes the double-counting that a plain OR-convolution would introduce.
  • What is OR-convolution vs AND-convolution vs XOR-convolution? OR-convolution = zeta (subset SOS) → pointwise multiply → Möbius. AND-convolution is the same with the superset transform. XOR-convolution needs the Walsh–Hadamard transform instead of zeta (cross-reference Phase 12 lab 07’s follow-up) — a different butterfly, but the same “transform, pointwise-multiply, inverse-transform” pipeline.
  • When does 3ⁿ beat 2ⁿ·n? When you need submask sums for only a few masks, or when the per-subset operation isn’t additive (so the incremental zeta relaxation doesn’t apply). For all-masks additive aggregation, 2ⁿ·n dominates.
  • How does SOS help other bitmask DPs? Precompute a table of “best value achievable by any submask of mask” via a max-flavored SOS (replace += with a max), then answer meet-in-the-middle or set-cover pruning queries in O(1). The transform generalizes to any commutative-monoid aggregate (sum, max, min, count).
  • How do you get superset sums instead of subset sums? Flip the merge direction: if not (mask & bit): F[mask] += F[mask | bit], or run subset SOS on the bit-complemented index space. Useful for AND-based counting where you want “supersets of a given mask.”

Product Extension

Feature-flag and permission systems: given a value assigned to each exact capability set, SOS answers “total value across all subsets of these permissions” in one pass — e.g. cohort counts over every attribute combination for analytics roll-ups. Ad/recommendation systems that bucket users by a bitmask of interests use subset/superset sums to answer “how many users have at least these interests” (superset) or “value summed over all interest subsets” (subset) without re-scanning. Any OLAP-style cube aggregation over a small set of boolean dimensions is an n-dimensional prefix sum — literally SOS.

Language/Runtime Follow-ups

  • Python: the tight 2ⁿ·n loop is slow in pure Python; for n = 20 prefer numpy vectorization per dimension (F[mask_has_bit] += F[mask_has_bit ^ bit] via boolean indexing) or accept the constant factor. Copy with A[:].
  • Java/C++: trivial for (i) for (mask) with int[]/vector<long long>; watch overflow — LC 982 counts reach 10⁹, use long/long long.
  • C++: this is the canonical competitive form; for(int i=0;i<n;i++) for(int m=0;m<(1<<n);m++) if(m>>i&1) F[m]+=F[m^1<<i];.
  • Go: int64 slice, same double loop; no generics needed.
  • JS/TS: Int32Array/Float64Array of length 2ⁿ; use Number up to 2⁵³, BigInt only if counts overflow.

Common Bugs

  1. Loop order swapped: iterating masks outer and dimensions inner computes garbage — the zeta relaxation requires each bit-dimension to fully sweep before the next.
  2. Missing the mask & bit guard: updating every mask (not just those with bit i set) double-adds and inflates sums.
  3. In-place clobber: transforming the caller’s array without copying corrupts data the caller still needs (e.g. reusing pairs after sos).
  4. Wrong universe size: sizing F to max(nums)+1 bits but then indexing ~z (which can set high bits) reads out of range or misses pairs. Mask complements with & FULL and size for the full bit width.
  5. Integer overflow: LC 982’s answer approaches m³ = 10⁹; a 32-bit accumulator overflows in Java/C/Go.
  6. Confusing subset and superset direction: using F[mask ^ bit] (subset) when the problem wants supersets (F[mask | bit]) inverts the meaning of the aggregate.

Debugging Strategy

Validate sos against the 3ⁿ subset loop on tiny n first — if they disagree, the transform itself is wrong, independent of the problem. Then check inverse_sos(sos(A)) == A to confirm the Möbius direction. For LC 982, run the O(m³) oracle on random small arrays: a mismatch localizes to either the pairs histogram (O(m²) step) or the complement lookup F[FULL ^ z]. Print F for a 2-bit universe by hand — F[0b11] should equal A[0]+A[1]+A[2]+A[3] — to catch loop-order and guard bugs instantly.

Mastery Criteria

  • Wrote the forward SOS transform from memory in <5 minutes.
  • Wrote the Möbius inverse and the 3ⁿ subset loop cold.
  • Explained SOS as an n-dimensional prefix sum with the loop invariant.
  • Solved LC 982 via the AND-to-submask reduction and SOS in <20 minutes.
  • Stated the crossover between 3ⁿ and 2ⁿ·n and picked correctly.
  • Connected zeta/Möbius to inclusion-exclusion and named the OR/AND/XOR-convolution pipeline (WHT for XOR).