Find K-Sum Subsets

Problem#

(curriculum variant.) Given an array nums and an integer k, return the k subsets with the largest sums (each as a list).

Examples#

nums = [1, 2, 3], k = 4 → [[1,2,3], [2,3], [1,3], [3]]

Constraints#

1 <= k <= 2^n.

Hints#

Hint 1

Sort descending. Max-heap holds (running sum, last-included index). Two children per pop: include next, or skip the included-last and include next.

Approach#

Sort nums descending. Max-heap of (sum, subset, lastIndex). Start with (0, [], -1) and the full subset is just adding the next available element. To avoid duplicates, restrict each pop’s expansions to indices > lastIndex; emit two branches: extend with next element, or swap last for next (Eppstein’s k-th shortest path expansion).

Solution#

vector<vector<int>> kLargestSubsetSums(vector<int> nums, int k) {
    sort(nums.rbegin(), nums.rend());
    using State = tuple<long long, vector<int>, int>; // (-sum for min-heap order, subset, lastIdx)
    priority_queue<State> pq;
    pq.push({0LL, {}, -1});
    vector<vector<int>> ans;
    while (k-- > 0 && !pq.empty()) {
        auto [sum, sub, last] = pq.top(); pq.pop();
        ans.push_back(sub);
        int nxt = last + 1;
        if (nxt < (int)nums.size()) {
            auto extended = sub; extended.push_back(nums[nxt]);
            pq.push({sum + nums[nxt], extended, nxt});
            if (!sub.empty()) {
                auto swapped = sub; swapped.back() = nums[nxt];
                pq.push({sum - nums[last] + nums[nxt], swapped, nxt});
            }
        }
    }
    return ans;
}

import (
    "container/heap"
    "sort"
)

type state struct {
    sum  int64
    sub  []int
    last int
}
type maxHeap []state

func (h maxHeap) Len() int            { return len(h) }
func (h maxHeap) Less(i, j int) bool  { return h[i].sum > h[j].sum }
func (h maxHeap) Swap(i, j int)       { h[i], h[j] = h[j], h[i] }
func (h *maxHeap) Push(x interface{}) { *h = append(*h, x.(state)) }
func (h *maxHeap) Pop() interface{} {
    old := *h
    n := len(old)
    x := old[n-1]
    *h = old[:n-1]
    return x
}

func kLargestSubsetSums(nums []int, k int) [][]int {
    sort.Sort(sort.Reverse(sort.IntSlice(nums)))
    pq := &maxHeap{}
    heap.Push(pq, state{0, []int{}, -1})
    ans := [][]int{}
    for k > 0 && pq.Len() > 0 {
        cur := heap.Pop(pq).(state)
        ans = append(ans, cur.sub)
        nxt := cur.last + 1
        if nxt < len(nums) {
            extended := make([]int, len(cur.sub), len(cur.sub)+1)
            copy(extended, cur.sub)
            extended = append(extended, nums[nxt])
            heap.Push(pq, state{cur.sum + int64(nums[nxt]), extended, nxt})
            if len(cur.sub) > 0 {
                swapped := make([]int, len(cur.sub))
                copy(swapped, cur.sub)
                swapped[len(swapped)-1] = nums[nxt]
                heap.Push(pq, state{cur.sum - int64(nums[cur.last]) + int64(nums[nxt]), swapped, nxt})
            }
        }
        k--
    }
    return ans
}

import heapq
from typing import List

def kLargestSubsetSums(nums: List[int], k: int) -> List[List[int]]:
    nums = sorted(nums, reverse=True)
    # Use negative sums for max-heap via heapq (min-heap)
    pq = [(0, [], -1)]  # (-sum, subset, lastIdx) -> but we want max, so store -sum
    # Re-init with -sum
    pq = [(-0, [], -1)]
    ans: List[List[int]] = []
    while k > 0 and pq:
        neg_sum, sub, last = heapq.heappop(pq)
        ans.append(sub)
        nxt = last + 1
        if nxt < len(nums):
            extended = sub + [nums[nxt]]
            heapq.heappush(pq, (neg_sum - nums[nxt], extended, nxt))
            if sub:
                swapped = sub[:-1] + [nums[nxt]]
                heapq.heappush(pq, (neg_sum + nums[last] - nums[nxt], swapped, nxt))
        k -= 1
    return ans

class MaxHeap {
    constructor() { this.h = []; }
    size() { return this.h.length; }
    push(x) {
        this.h.push(x);
        let i = this.h.length - 1;
        while (i > 0) {
            const p = (i - 1) >> 1;
            if (this.h[i].sum > this.h[p].sum) {
                [this.h[i], this.h[p]] = [this.h[p], this.h[i]];
                i = p;
            } else break;
        }
    }
    pop() {
        const top = this.h[0];
        const last = this.h.pop();
        if (this.h.length > 0) {
            this.h[0] = last;
            let i = 0;
            const n = this.h.length;
            while (true) {
                const l = i * 2 + 1, r = i * 2 + 2;
                let best = i;
                if (l < n && this.h[l].sum > this.h[best].sum) best = l;
                if (r < n && this.h[r].sum > this.h[best].sum) best = r;
                if (best === i) break;
                [this.h[i], this.h[best]] = [this.h[best], this.h[i]];
                i = best;
            }
        }
        return top;
    }
}

function kLargestSubsetSums(nums, k) {
    nums = [...nums].sort((a, b) => b - a);
    const pq = new MaxHeap();
    pq.push({ sum: 0, sub: [], last: -1 });
    const ans = [];
    while (k > 0 && pq.size() > 0) {
        const cur = pq.pop();
        ans.push(cur.sub);
        const nxt = cur.last + 1;
        if (nxt < nums.length) {
            pq.push({ sum: cur.sum + nums[nxt], sub: [...cur.sub, nums[nxt]], last: nxt });
            if (cur.sub.length > 0) {
                const swapped = cur.sub.slice(0, -1);
                swapped.push(nums[nxt]);
                pq.push({ sum: cur.sum - nums[cur.last] + nums[nxt], sub: swapped, last: nxt });
            }
        }
        k--;
    }
    return ans;
}

class Solution {
    public List<List<Integer>> kLargestSubsetSums(int[] numsIn, int k) {
        Integer[] boxed = new Integer[numsIn.length];
        for (int i = 0; i < numsIn.length; i++) boxed[i] = numsIn[i];
        Arrays.sort(boxed, Comparator.reverseOrder());
        int[] nums = new int[boxed.length];
        for (int i = 0; i < boxed.length; i++) nums[i] = boxed[i];

        PriorityQueue<State> pq = new PriorityQueue<>((a, b) -> Long.compare(b.sum, a.sum));
        pq.offer(new State(0L, new ArrayList<>(), -1));
        List<List<Integer>> ans = new ArrayList<>();
        while (k > 0 && !pq.isEmpty()) {
            State cur = pq.poll();
            ans.add(cur.sub);
            int nxt = cur.last + 1;
            if (nxt < nums.length) {
                List<Integer> extended = new ArrayList<>(cur.sub);
                extended.add(nums[nxt]);
                pq.offer(new State(cur.sum + nums[nxt], extended, nxt));
                if (!cur.sub.isEmpty()) {
                    List<Integer> swapped = new ArrayList<>(cur.sub);
                    swapped.set(swapped.size() - 1, nums[nxt]);
                    pq.offer(new State(cur.sum - nums[cur.last] + nums[nxt], swapped, nxt));
                }
            }
            k--;
        }
        return ans;
    }

    private static class State {
        long sum;
        List<Integer> sub;
        int last;
        State(long s, List<Integer> sub, int last) {
            this.sum = s; this.sub = sub; this.last = last;
        }
    }
}

interface SubsetState {
    sum: number;
    sub: number[];
    last: number;
}

class MaxHeap {
    private h: SubsetState[] = [];
    size(): number { return this.h.length; }
    push(x: SubsetState): void {
        this.h.push(x);
        let i = this.h.length - 1;
        while (i > 0) {
            const p = (i - 1) >> 1;
            if (this.h[i].sum > this.h[p].sum) {
                [this.h[i], this.h[p]] = [this.h[p], this.h[i]];
                i = p;
            } else break;
        }
    }
    pop(): SubsetState {
        const top = this.h[0];
        const last = this.h.pop()!;
        if (this.h.length > 0) {
            this.h[0] = last;
            let i = 0;
            const n = this.h.length;
            while (true) {
                const l = i * 2 + 1, r = i * 2 + 2;
                let best = i;
                if (l < n && this.h[l].sum > this.h[best].sum) best = l;
                if (r < n && this.h[r].sum > this.h[best].sum) best = r;
                if (best === i) break;
                [this.h[i], this.h[best]] = [this.h[best], this.h[i]];
                i = best;
            }
        }
        return top;
    }
}

function kLargestSubsetSums(nums: number[], k: number): number[][] {
    const sorted = [...nums].sort((a, b) => b - a);
    const pq = new MaxHeap();
    pq.push({ sum: 0, sub: [], last: -1 });
    const ans: number[][] = [];
    while (k > 0 && pq.size() > 0) {
        const cur = pq.pop();
        ans.push(cur.sub);
        const nxt = cur.last + 1;
        if (nxt < sorted.length) {
            pq.push({ sum: cur.sum + sorted[nxt], sub: [...cur.sub, sorted[nxt]], last: nxt });
            if (cur.sub.length > 0) {
                const swapped = cur.sub.slice(0, -1);
                swapped.push(sorted[nxt]);
                pq.push({ sum: cur.sum - sorted[cur.last] + sorted[nxt], sub: swapped, last: nxt });
            }
        }
        k--;
    }
    return ans;
}

Editorial#

Identical structure to Find K-Sum of an Array. Sorting descending ensures that “extend with the next element” or “swap last for next” both produce sums that decrease — so heap-pop order yields subsets in descending sum order. Each subset is emitted exactly once via the no-overlap invariant on lastIndex.

Complexity#

Time: O((n + k) log (n + k) * n).
Space: O((n + k) * n) (subsets stored).

Problem#

Examples#

Constraints#

Hints#

Approach#

Solution#

Editorial#

Complexity#

Concept revision#

Keyboard shortcuts

Problem#

Examples#

Constraints#

Hints#

Approach#

Solution#

Editorial#

Complexity#

Concept revision#

Related#