K Maximum Sum Combinations From Two Arrays

vector<int> kMaxSumCombinations(vector<int> a, vector<int> b, int k) {
    sort(a.rbegin(), a.rend());
    sort(b.rbegin(), b.rend());
    priority_queue<tuple<int,int,int>> pq;
    set<pair<int,int>> seen;
    pq.emplace(a[0] + b[0], 0, 0);
    seen.insert({0, 0});
    vector<int> ans;
    while ((int)ans.size() < k && !pq.empty()) {
        auto [s, i, j] = pq.top(); pq.pop();
        ans.push_back(s);
        if (i + 1 < (int)a.size() && !seen.count({i + 1, j})) {
            pq.emplace(a[i + 1] + b[j], i + 1, j);
            seen.insert({i + 1, j});
        }
        if (j + 1 < (int)b.size() && !seen.count({i, j + 1})) {
            pq.emplace(a[i] + b[j + 1], i, j + 1);
            seen.insert({i, j + 1});
        }
    }
    return ans;
}

import (
    "container/heap"
    "sort"
)

type tuple struct{ s, i, j int }
type maxTupleHeap []tuple

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

func kMaxSumCombinations(a, b []int, k int) []int {
    sort.Sort(sort.Reverse(sort.IntSlice(a)))
    sort.Sort(sort.Reverse(sort.IntSlice(b)))
    pq := &maxTupleHeap{}
    seen := map[[2]int]bool{}
    heap.Push(pq, tuple{a[0] + b[0], 0, 0})
    seen[[2]int{0, 0}] = true
    ans := []int{}
    for len(ans) < k && pq.Len() > 0 {
        t := heap.Pop(pq).(tuple)
        ans = append(ans, t.s)
        if t.i+1 < len(a) && !seen[[2]int{t.i + 1, t.j}] {
            heap.Push(pq, tuple{a[t.i+1] + b[t.j], t.i + 1, t.j})
            seen[[2]int{t.i + 1, t.j}] = true
        }
        if t.j+1 < len(b) && !seen[[2]int{t.i, t.j + 1}] {
            heap.Push(pq, tuple{a[t.i] + b[t.j+1], t.i, t.j + 1})
            seen[[2]int{t.i, t.j + 1}] = true
        }
    }
    return ans
}

from typing import List
import heapq

class Solution:
    def kMaxSumCombinations(self, a: List[int], b: List[int], k: int) -> List[int]:
        a = sorted(a, reverse=True)
        b = sorted(b, reverse=True)
        pq: List = [(-(a[0] + b[0]), 0, 0)]
        seen = {(0, 0)}
        ans: List[int] = []
        while len(ans) < k and pq:
            negs, i, j = heapq.heappop(pq)
            ans.append(-negs)
            if i + 1 < len(a) and (i + 1, j) not in seen:
                heapq.heappush(pq, (-(a[i + 1] + b[j]), i + 1, j))
                seen.add((i + 1, j))
            if j + 1 < len(b) and (i, j + 1) not in seen:
                heapq.heappush(pq, (-(a[i] + b[j + 1]), i, j + 1))
                seen.add((i, j + 1))
        return ans

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

function kMaxSumCombinations(a, b, k) {
    a = [...a].sort((x, y) => y - x);
    b = [...b].sort((x, y) => y - x);
    const pq = new MaxHeap();
    const seen = new Set();
    pq.push([a[0] + b[0], 0, 0]);
    seen.add('0,0');
    const ans = [];
    while (ans.length < k && pq.size()) {
        const [s, i, j] = pq.pop();
        ans.push(s);
        const k1 = `${i + 1},${j}`;
        if (i + 1 < a.length && !seen.has(k1)) {
            pq.push([a[i + 1] + b[j], i + 1, j]);
            seen.add(k1);
        }
        const k2 = `${i},${j + 1}`;
        if (j + 1 < b.length && !seen.has(k2)) {
            pq.push([a[i] + b[j + 1], i, j + 1]);
            seen.add(k2);
        }
    }
    return ans;
}

class Solution {
    public List<Integer> kMaxSumCombinations(int[] aIn, int[] bIn, int k) {
        Integer[] a = new Integer[aIn.length];
        Integer[] b = new Integer[bIn.length];
        for (int i = 0; i < aIn.length; i++) a[i] = aIn[i];
        for (int i = 0; i < bIn.length; i++) b[i] = bIn[i];
        Arrays.sort(a, Comparator.reverseOrder());
        Arrays.sort(b, Comparator.reverseOrder());
        PriorityQueue<int[]> pq = new PriorityQueue<>((x, y) -> y[0] - x[0]);
        Set<Long> seen = new HashSet<>();
        pq.offer(new int[]{a[0] + b[0], 0, 0});
        seen.add(0L);
        List<Integer> ans = new ArrayList<>();
        while (ans.size() < k && !pq.isEmpty()) {
            int[] t = pq.poll();
            int s = t[0], i = t[1], j = t[2];
            ans.add(s);
            if (i + 1 < a.length) {
                long key = (long)(i + 1) * 100000L + j;
                if (seen.add(key)) {
                    pq.offer(new int[]{a[i + 1] + b[j], i + 1, j});
                }
            }
            if (j + 1 < b.length) {
                long key = (long) i * 100000L + (j + 1);
                if (seen.add(key)) {
                    pq.offer(new int[]{a[i] + b[j + 1], i, j + 1});
                }
            }
        }
        return ans;
    }
}

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

function kMaxSumCombinations(a: number[], b: number[], k: number): number[] {
    const sortedA = [...a].sort((x, y) => y - x);
    const sortedB = [...b].sort((x, y) => y - x);
    const pq = new MaxHeap();
    const seen: Set<string> = new Set();
    pq.push([sortedA[0] + sortedB[0], 0, 0]);
    seen.add('0,0');
    const ans: number[] = [];
    while (ans.length < k && pq.size()) {
        const [s, i, j] = pq.pop();
        ans.push(s);
        const k1 = `${i + 1},${j}`;
        if (i + 1 < sortedA.length && !seen.has(k1)) {
            pq.push([sortedA[i + 1] + sortedB[j], i + 1, j]);
            seen.add(k1);
        }
        const k2 = `${i},${j + 1}`;
        if (j + 1 < sortedB.length && !seen.has(k2)) {
            pq.push([sortedA[i] + sortedB[j + 1], i, j + 1]);
            seen.add(k2);
        }
    }
    return ans;
}