# Find k smallest sum n-tuples

Given a collection of sorted lists {l1, l2, ...} I need to find the smallest k index tuples taken from these lists by summed value, e.g. given:

{
{1, 2, 3},
{5, 6, 7},
{3, 4, 5}
}


If k were 3 I'd want:

{
{1, 1, 1},
{2, 1, 1},
{1, 1, 2}
}


or

{
{1, 1, 1},
{2, 1, 1},
{3, 1, 1}
}


where these are indices into the original lists giving the smallest sums. This is basically what I'm talking about.

As Henrik Schumacher points out, the question of what to do with ties is unspecified and I have no reason to pick a given way to resolve them, so let's let that be dealt with however one may want.

The question now is how to find this when the k, the number of lists, and the length of these lists get large.

I've implemented a compiled solution based on this but I want to know if a) there's something built in or b) there's a better way to do it.

Here's my code as it stands:

getTupleSetsC =
Compile[
{{vals, _Real, 2}, {n, _Integer}},
Module[
{
dim,
queue,
costs,
defel,
tups,
pop,
push,
cost,
qels = 0,
solns = 0,
order
},
dim = Length@vals;
(* Initialize priority queue and cost vector *)

defel = Table[-1, {$, dim}]; queue = Table[defel, {$, n}];
costs = Table[-1., {$, n}]; (* Initialize solutions *) tups = Table[defel, {$, n}];
(* Determine first element *)
push = Table[1, {$, dim}]; (* Compute cost *) cost = 0.; Do[cost += vals[[i, push[[i]]]], {i, dim}]; (* Push onto queue *) ++qels; queue[[qels]] = push; costs[[qels]] = cost; (* Iterate *) pop = defel; Do[ (* pop the first element, add it to the tups array *) pop = queue[]; tups[[++solns]] = pop; (* shift the queue elements back *) queue[[;; qels]] = queue[[2 ;; qels + 1]]; costs[[;; qels]] = costs[[2 ;; qels + 1]]; qels--; Do[ (* Create new push element *) push = pop; push[[j]] += 1; (* Compute cost *) cost = 0.; Do[cost += vals[[i, push[[i]]]], {i, dim}]; (* determine where on queue would be added *) ++qels; (* allocate more queue if necessary *) If[qels > Length@queue, queue = Join[queue, Table[defel, {$$, 2*Length@queue}] ]; costs = Join[costs, Table[-1., {$$, 2*Length@queue}] ] ]; (* push onto queue *) queue[[qels]] = push; costs[[qels]] = cost; (* If next element would be added in later iteration, break *) If[pop[[j]] > 1, Break[]], {j, dim} ]; (* resort by priority *) order = Ordering[costs[[;; qels]]]; queue[[;; qels]] = queue[[;; qels]][[order]]; costs[[;; qels]] = costs[[;; qels]][[order]], {$, n - 1}
];
pop = queue[];
tups[[++solns]] = pop;
tups
]
];

getTupleSets[lists_?(MatrixQ[#, InternalRealValuedNumericQ] &),
n_Integer?Positive] :=
Module[
{
pos,
sets
},
pos = getTupleSetsC[lists, n];
sets = MapIndexed[lists[[#2[], #]] &, Transpose@pos] // Transpose;
]

StringFreeQ[CompilePrint[getTupleSetsC], "MainEvaluate"]

True


And it's fast:

n = 100;
d = 6;
tests =
Sort /@
BlockRandom[
SeedRandom;
Table[RandomSample[Range[0, 500], n], d]
];

getTupleSets[tests, 5] // RepeatedTiming

{0.000095, {{1, 1, 1, 1, 1, 1} -> {7, 2, 7, 0, 14, 8}, {1, 1, 1, 2, 1,
1} -> {7, 2, 7, 1, 14, 8}, {2, 1, 1, 1, 1, 1} -> {9, 2, 7, 0, 14,
8}, {1, 1, 1, 1, 2, 1} -> {7, 2, 7, 0, 17, 8}, {1, 1, 1, 1, 1,
2} -> {7, 2, 7, 0, 14, 11}}}

getTupleSets[tests, n]; // RepeatedTiming

{0.0013, Null}


But it's got a few issues that might cause serious slow downs and which I don't know if I want to optimize away until I know if there's a better approach. For one, I use Sort a full k times on data of length up-to around k. This could potentially be inefficient. I also potentially use too much memory, as I don't know if it'd be more efficient to start with less.

Can anyone help me do better?

• What is your measure of smallness? You don't mean the lexicographical order, do you? I don't quite get why you consider {7, 2, 7, 0, 17, 8} smaller than {7, 2, 7, 0, 14, 11}. Jul 31 '18 at 5:20
• @HenrikSchumacher ah... I don't really know what should happen in the case of ties. That tie was simply broken by where the 2 appeared in the index list. I just meant smallness by summation (I'm using this to get tuples of eigenvalues with smallest sum). Jul 31 '18 at 5:22
• So {1,1,2} differs from {1,2,1}? Jul 31 '18 at 6:06
• @DavidG.Stork yeah maybe I wasn't clear enough. Those are indices into the original list. It's about picking the k tuples of indices that yield the smallest sum when applied to the original lists. Jul 31 '18 at 6:07
• Hm. I guess, I would implement a breadth-first search. Skimming your code, that's what you have already done, right? Jul 31 '18 at 7:21

# Bug Fixes

I believe there are two bugs in getTupleSetsC implementation. When shifting queue elements back, part specification: 2 ;; qels + 1 is used before decrementing qels which causes Part::take error when queue is full. Also when creating new push element push[[j]] += 1 is done unconditionally which leads to unbounded increase of indices. After fixing those bugs we get function which we'll use as reference for benchmarks:

getTupleSetsCb3m2a1 =
Compile[
{{vals, _Real, 2}, {n, _Integer}},
Module[
{
dim,
k,
queue,
costs,
defel,
tups,
pop,
push,
cost,
qels = 0,
solns = 0,
order
},
(*dim = Length@vals;*)
{dim, k} = Dimensions@vals;
(* Initialize priority queue and cost vector *)

defel = Table[-1, {$$, dim}]; queue = Table[defel, {, n}]; costs = Table[-1., {, n}]; (* Initialize solutions *) tups = Table[defel, {, n}]; (* Determine first element *) push = Table[1, {$$, dim}];
(* Compute cost *)
cost = 0.;
Do[cost += vals[[i, push[[i]]]], {i, dim}];
(* Push onto queue *)
++qels;
queue[[qels]] = push;
costs[[qels]] = cost;
(* Iterate *)
pop = defel;
Do[
(* pop the first element, add it to the tups array *)

pop = queue[];
tups[[++solns]] = pop;
(* shift the queue elements back *)

(*queue[[;; qels]] = queue[[2 ;; qels + 1]];
costs[[;; qels]] = costs[[2 ;; qels + 1]];
qels--;*)
qels--;
If[qels > 0,
queue[[;; qels]] = queue[[2 ;; qels + 1]];
costs[[;; qels]] = costs[[2 ;; qels + 1]];
];
Do[
(* If j-th index is already maximal there's no tuple with incremented j-th index.
Since it also means that j-th index > 1, we can not only continue to next j, but break the loop. *)
If[pop[[j]] == k, Break[]];

(* Create new push element *)
push = pop;
push[[j]] += 1;
(* Compute cost *)
cost = 0.;
Do[cost += vals[[i, push[[i]]]], {i, dim}];
(* determine where on queue would be added *)
++qels;
(* allocate more queue if necessary *)

If[qels > Length@queue,
queue =
Join[queue,
Table[defel, {$$, 2*Length@queue}] ]; costs = Join[costs, Table[-1., {$$, 2*Length@queue}]
]
];
(* push onto queue *)
queue[[qels]] = push;
costs[[qels]] = cost;
(* If next element would be added in later iteration, break *)

If[pop[[j]] > 1, Break[]],
{j, dim}
];
(* resort by priority *)

order = Ordering[costs[[;; qels]]];
queue[[;; qels]] = queue[[;; qels]][[order]];
costs[[;; qels]] = costs[[;; qels]][[order]],
{\$, n - 1}
];
pop = queue[];
tups[[++solns]] = pop;
tups
],
CompilationTarget -> "C",
RuntimeOptions -> "Speed"
] // Last;


# Performance improvements

I don't have better algorithm than one used by OP. So let's just change few implementation details.

First obvious, and already hinted by OP, thing is priority queue implementation. Instead of shifted and resorted array let's use heap-ordered binary tree.

As explained in linked answer to CSSE question our queue can be bounded, and when maximal size is reached we'll remove unneeded most expensive elements. Heap is not an asymptotically optimal data structure when comes to removing most expensive elements, but in this particular algorithm, I believe, our heap operates almost all the time below its maximal capacity. So occasional, if present at all, non-optimal removals will hopefully be "amortized". And we'll benefit from lower constants in optimal operations of adding arbitrary elements and removing cheapest ones.

To not waste time on swapping whole tuples when restoring heap property, we'll use heap of indices to base arrays containing tuples and their costs. To quickly find empty spots in those base arrays we'll also keep an array based linked list of empty spots.

We'll keep track of largest cost of tuples present in queue and won't push a tuple if we already have enough cheaper ones.

Since new tuple is always created from old one by changing one index, we don't have to loop over all its elements to calculate its cost, it's enough to focus on changed element: subtract cost of old element and add cost of new one.

Above changes result in following code:

getTupleSetsCjkuczm = Hold@Compile[{{vals, _Real, 2}, {k, _Integer}},
Module[{m, n, d, firstEmpty, noEmpty, nextEmpty, queueCurrent, queue, currentCost, largestCost, emptyCost, costs, tuples, result},
{m, n} = Dimensions@vals;
d = k + m;
firstEmpty = 1;
noEmpty = d + 1;
nextEmpty = Range[2, noEmpty];
queueCurrent = 1;
queue = Table[0, d];
largestCost = 0.;
emptyCost = 2. m Max@vals;
costs = Table[emptyCost, d];
tuples = Table[1, d, m];
result = Table[1, k, m];

(* Compute cost of first tuple. *)
currentCost = vals[[1, 1]];
Do[currentCost += vals[[i, 1]], {i, 2, m}];

Do[
Do[
Module[{oldInd, newInd, newCost},
oldInd = result[[i, j]];

(* If j-th index is already maximal there's no tuple with incremented j-th index.
Since it also means that j-th index > 1, we can not only continue to next j, but break the loop. *)
If[oldInd == n, Break[]];

newInd = oldInd + 1;
newCost = currentCost - vals[[j, oldInd]] + vals[[j, newInd]];

(* If we already have enough cheaper tuples, don't push this one. *)
If[queueCurrent > k - i && newCost >= largestCost, If[oldInd > 1, Break[], Continue[]]];

(* If there are no empty spots, leave only k - i cheapes elements, we don't need rest. *)
If[firstEmpty == noEmpty,
Module[{oPrev},
(* Create new queue. This is slow, but should be very rare. *)
queueCurrent = k - i + 1;
queue = Ordering@costs;
largestCost = costs[[queue[[k - i]]]];

(* Update list of empty spots. *)
oPrev = queue[[queueCurrent]];
firstEmpty = oPrev;
Do[
With[{o = queue[[l]]},
nextEmpty[[oPrev]] = o;
oPrev = o;
]
,
{l, k - i + 2, d}
];
nextEmpty[[oPrev]] = noEmpty;
];
];

(* Store tuple and its cost. *)
tuples[[firstEmpty]] = result[[i]];
tuples[[firstEmpty, j]] = newInd;
costs[[firstEmpty]] = newCost;
If[newCost > largestCost, largestCost = newCost];

(* Push index into queue. *)
queue[[queueCurrent++]] = firstEmpty;
(* Restore heap property. *)
Module[{l, p},
l = queueCurrent - 1;
p = Quotient[l - 2, 2] + 1;
While[l > 1 && costs[[queue[[l]]]] < costs[[queue[[p]]]],
With[{lth = queue[[l]]},
queue[[l]] = queue[[p]];
queue[[p]] = lth;
];
l = p;
p = Quotient[l - 2, 2] + 1;
];
];

(* Remove first empty spot from list of empty spots. *)
firstEmpty = nextEmpty[[firstEmpty]];

(* If next element would be added in later iteration, break. *)
If[oldInd > 1, Break[]]
]
,
{j, m}
];

Module[{cheapest},
(* Pop cheapest element from queue. *)
cheapest = queue[];
queue[] = queue[[--queueCurrent]];
(* Restore heap property. *)
Module[{j = 1},
While[True,
Module[{l, r, o, jth, jthCost},
l =  2 j;
r = 2 j + 1;
jth = queue[[j]];
jthCost = costs[[jth]];
o = If[r < queueCurrent && costs[[queue[[r]]]] < jthCost,
If[costs[[queue[[l]]]] < costs[[queue[[r]]]], l, r]
(* else *),
If[l < queueCurrent && costs[[queue[[l]]]] < jthCost, l, 0]
];
If[o < 1, Break[]];
queue[[j]] = queue[[o]];
queue[[o]] = jth;
j = o;
];
];
];

currentCost = costs[[cheapest]];
costs[[cheapest]] = emptyCost;

(* Add cheapest element to result array. *)
result[[i + 1]] = tuples[[cheapest]];

(* Prepend cheapest index to list of empty indices. *)
nextEmpty[[cheapest]] = firstEmpty;
firstEmpty = cheapest;
];
,
{i, k - 1}
];
result
],
CompilationTarget -> "C",
RuntimeOptions -> "Speed"
] /. Part -> CompileGetElement //. HoldPattern[CompileGetElement[x__] = y_] :> (Part@x = y) // ReleaseHold // Last;


# Tests

On small data sets we can test that we get same results as ones obtained from Sorted lists of all Tuples:

getCosts = Total[Transpose@MapThread[#1[[#2]] &, {#1, Transpose@#2}], {2}] &;

SeedRandom@0
testData = Sort /@ RandomInteger[{0, 500}, {2, 4}]
res = Sort@Total[Tuples@testData, {2}];
Array[getCosts[testData, getTupleSetsCb3m2a1[testData, #]] === Take[res, #] &, Length@res, 1, And] // AbsoluteTiming
Array[getCosts[testData, getTupleSetsCjkuczm[testData, #]] === Take[res, #] &, Length@res, 1, And] // AbsoluteTiming
(* {{32, 266, 334, 399}, {30, 270, 272, 414}} *)
(* {0.00083,  True} *)
(* {0.000243, True} *)

testData = Sort /@ RandomInteger[{0, 500}, {6, 3}]
res = Sort@Total[Tuples@testData, {2}];
Array[getCosts[testData, getTupleSetsCb3m2a1[testData, #]] === Take[res, #] &, Length@res, 1, And] // AbsoluteTiming
Array[getCosts[testData, getTupleSetsCjkuczm[testData, #]] === Take[res, #] &, Length@res, 1, And] // AbsoluteTiming
(* {{195, 266, 342}, {138, 309, 473}, {264, 338, 351}, {111, 126, 356}, {1, 77, 179}, {210, 421, 495}} *)
(* {1.44529,  True} *)
(* {0.089297, True} *)

testData = Sort /@ RandomInteger[{0, 500}, {5, 5}]
res = Sort@Total[Tuples@testData, {2}];
Array[getCosts[testData, getTupleSetsCb3m2a1[testData, #]] === Take[res, #] &, Length@res, 1, And] // AbsoluteTiming
Array[getCosts[testData, getTupleSetsCjkuczm[testData, #]] === Take[res, #] &, Length@res, 1, And] // AbsoluteTiming
(* {{14, 201, 249, 315, 487}, {269, 321, 340, 411, 464}, {23, 229, 357, 372, 377}, {89, 205, 297, 420, 442}, {171, 268, 298, 374, 481}} *)
(* {71.903,   True} *)
(* { 1.43877, True} *)


# Benchmarks

Let's compare time and memory usage on data from OP:

tests = Sort /@ BlockRandom[SeedRandom@1; Table[RandomSample[Range[0, 500], 100], 6]];
res1 = getTupleSetsCb3m2a1[tests, 100]; // MaxMemoryUsed // RepeatedTiming
res2 = getTupleSetsCjkuczm[tests, 100]; // MaxMemoryUsed // RepeatedTiming
getCosts[tests, res1] === getCosts[tests, res2]
(* {0.000523, 59296} *)
(* {0.000032, 24032} *)
(* True *)

res1 = getTupleSetsCb3m2a1[tests, 10^4]; // MaxMemoryUsed // RepeatedTiming
res2 = getTupleSetsCjkuczm[tests, 10^4]; // MaxMemoryUsed // RepeatedTiming
getCosts[tests, res1] === getCosts[tests, res2]
(* {1.24,    1783744} *)
(* {0.00288, 1673128} *)
(* True *)


And on other combinations of data sizes:

Needs@"GeneralUtilities"
Table[
With[{testData = Sort /@ RandomInteger[{0, 500}, {i, j}]},
BenchmarkPlot[<|"b3m2a1" -> (getTupleSetsCb3m2a1[testData, #] &), "jkuczm" -> (getTupleSetsCjkuczm[testData, #] &)|>, Identity, 2^Range@16]
]
,
{i, 2^(3 Range@3)}, {j, 2^(3 Range@3)}
] // TableForm[#, TableHeadings -> {2^(3 Range@3), 2^(3 Range@3)}, TableAlignments -> Center] & • I'm thoroughly impressed each time I read this. I'll have to come back to this if I need a good compilable heap in the future. Oct 5 '18 at 6:46

Could do this as a polynomial algebra problem with matrix values used as exponents, with a distinct variable for each, summed in rows and product taken of row sums. Use a distinct variable to record total degrees. We set this up as below.

mat = {{1, 2, 3}, {5, 6, 7}, {3, 4, 5}};
xmat = Array[x, Dimensions[mat]];
prod = Apply[Times, Apply[Plus, xmat^mat, {1}]] /.

(* Out= (t x[1, 1] + t^2 x[1, 2]^2 + t^3 x[1, 3]^3) (t^5 x[2, 1]^5 +
t^6 x[2, 2]^6 + t^7 x[2, 3]^7) (t^3 x[3, 1]^3 + t^4 x[3, 2]^4 +
t^5 x[3, 3]^5) *)


Extract the min possible value, as the minimal exponent in t.

minexpon = Exponent[prod, t, Min]

(* Out= 9 *)


Now we define a function to extract lists of matrix column indices that all give rise to a particular power (row being denoted by position in the list). This version also gives the matrix entries in a separate list

getIndexVectors[prod_, deg_] := Module[
{coeffs},
coeffs = Coefficient[prod, t, deg];
If [Head[coeffs] === Plus, coeffs = Apply[List, coeffs],
coeffs = {coeffs}];
Map[Transpose, coeffs /. Times -> List /. x[i_, j_]^k_. :> {j, k}]
]


Some tests:

In:= getIndexVectors[prod, minexpon]

(* Out= {{{1, 1, 1}, {1, 5, 3}}} *)

In:= getIndexVectors[prod, minexpon + 1]

(* Out= {{{2, 1, 1}, {2, 5, 3}}, {{1, 2, 1}, {1, 6, 3}}, {{1, 1,
2}, {1, 5, 4}}} *)

• That's an interesting idea that I'd need to sit down and think about to understand I think. How fast do you think this can be? Speed and memory are my primary concerns here. Aug 30 '18 at 18:37
• If does not seem to mind long vectors, but if there are more than a few of them it gets quite slow. Might be better to use a greedy dynamic heuristic: start with position of smallest elements, then compare with all but one smallest, and the remaining one second smallest in its row, etc. Aug 30 '18 at 21:48
k = 5;
mylist = {{1, 2, 3}, {5, 6, 7}, {3, 4, 5}};
Take[
DeleteDuplicates[
Map[Sort,
SortBy[Tuples[DeleteDuplicates[Flatten[mylist]],
Length[First[mylist]]], Total]]], k]

• I think perhaps I was unclear about how these should be selected. They're selected by minimizing the sum of the elements. I'll make the example clearer. Meantime, this unfortunately does not do what I need. It could be trivially adapted to do so, but the number of Tuples will quickly destroy my memory. Jul 31 '18 at 5:57