4
$\begingroup$

Suppose you have the list $a=(k_1,k_2,k_3, \dots, k_N)$ such that $k_{i}\; (i=1,\ldots, N)$ are non negative integers. How can you get all the lists that satisfy the constraint $\sum_{i=1}^{N} k_{i}=M$?

$\endgroup$
2
  • 2
    $\begingroup$ Take a look at IntegerPartitions. $\endgroup$ – JimB Jan 19 at 17:31
  • 3
    $\begingroup$ ...or FrobeniusSolve[]. $\endgroup$ – J. M.'s ennui Jan 19 at 17:32
4
$\begingroup$

You can do it with FrobeniusSolve like so:

ClearAll[solutions]
byFrobenius[m_Integer?Positive, n_Integer?Positive] :=
  FrobeniusSolve[ConstantArray[1, n], m];

This works like so:

byFrobenius[5, 3]
(* {{0, 0, 5}, {0, 1, 4}, {0, 2, 3}, 
    {0, 3, 2}, {0, 4, 1}, {0, 5, 0}, 
    {1, 0, 4}, {1, 1, 3}, {1, 2, 2}, 
    {1, 3, 1}, {1, 4, 0}, {2, 0, 3}, 
    {2, 1, 2}, {2, 2, 1}, {2, 3, 0}, 
    {3, 0, 2}, {3, 1, 1}, {3, 2, 0}, 
    {4, 0, 1}, {4, 1, 0}, {5, 0, 0}} *) 

Let's compare this to a brute force solution:

ClearAll[byBruteForce];
byBruteForce[m_Integer?Positive, n_Integer?Positive] :=
  With[{candidates = Tuples[Range[0, m], n]},
   Select[candidates, Total /* EqualTo[m]]];
ContainsExactly[byBruteForce[5, 3], byFrobenius[5, 3]]
(* True *)

Performance of byFrobenius is, of course, much better:

byBruteForce[10, 6]; // AbsoluteTiming
(* {2.64395, Null} *)

byFrobenius[10, 6]; // AbsoluteTiming
(* {0.0364338, Null} *)

We can get even faster by using IntegerPartitions, as suggested by @geom, but remember to use the second argument to ensure we don't get lists that are too long (which will then be truncated by PadRight).

byIntegerPartitions[m_Integer?Positive, n_Integer?Positive] := 
 Catenate[Permutations /@ 
   PadRight[IntegerPartitions[m, n], {Automatic, n}]]

This is fast and correct:

byIntegerPartitions[10, 6]; // AbsoluteTiming
(* {0.0007277, Null} *)

ContainsExactly[byIntegerPartitions[10, 6], byFrobenius[10, 6]]
(* True *)

Interestingly, if we do some memoization, a hand-written, recursive solution can beat byFrobenius for performance (though it's still much slower than byIntegerPartitions):

byRecursion[m_Integer?Positive, n_Integer?Positive] :=
  Module[{memo},
   memo[k_, 1] := {{k}};
   memo[0, l_] := {ConstantArray[0, l]};
   memo[k_, l_] := memo[k, l] =
     Catenate[Map[Map[Prepend[#], memo[k - #, l - 1]] &, Range[0, k]]];
   
   memo[m, n]];

byRecursion[10, 6]; // AbsoluteTiming
(* {0.0172348, Null} *)

ContainsExactly[byRecursion[10, 6], byFrobenius[10, 6]]
(* True *)
$\endgroup$
1
  • $\begingroup$ Thanks for the detailed answer! $\endgroup$ – geom Jan 19 at 19:27
3
$\begingroup$

I think the answer is:

f[m_,n_]:=Flatten[Permutations /@ PadRight[IntegerPartitions[m], {Automatic, n}], 1]

e.g

f[2, 4]
(*=> {{2, 0, 0, 0}, {0, 2, 0, 0}, {0, 0, 2, 0}, {0, 0, 0, 2}, {1, 1, 0, 0}, 
{1, 0, 1, 0}, {1, 0, 0, 1}, {0, 1, 1, 0}, {0, 1, 0, 1}, {0, 0, 1, 1}} *)

An answer with FrobeniusSolve[] as suggested by @J.M.'s ennui, or any improvements in my answer would be appreciated

EDIT After @Pillsy's comment my answer should be modified as:

f[m_,n_]:=Catenate[Permutations /@ PadRight[IntegerPartitions[m,n], {Automatic, n}]]
$\endgroup$
2
  • 1
    $\begingroup$ This contains results that don't satisfy the constraints. Length@Select[f[10, 6], Total /* UnequalTo[10]] gives 84. $\endgroup$ – Pillsy Jan 19 at 19:12
  • $\begingroup$ @Pillsy True! I missed it. I modified my answer $\endgroup$ – geom Jan 19 at 19:46
3
$\begingroup$
f1[n_, m_] := IntegerPartitions[m, {n}, Range[m, 0, -1]]
f2[n_, m_] := 
 FrobeniusSolve[ConstantArray[1, n], m] // DeleteDuplicatesBy[Sort]

(Please notice that $N$ is before $M$) :P

Validation

f1[6, 10] // RepeatedTiming
f2[6, 10] // RepeatedTiming
%[[2]] === %%[[2]]
{8.628*10^-6, {{0, 0, 0, 0, 0, 10}, {0, 0, 0, 0, 1, 9}, {0, 0, 0, 0, 
   2, 8}, {0, 0, 0, 0, 3, 7}, {0, 0, 0, 0, 4, 6}, {0, 0, 0, 0, 5, 
   5}, {0, 0, 0, 1, 1, 8}, {0, 0, 0, 1, 2, 7}, {0, 0, 0, 1, 3, 6}, {0,
    0, 0, 1, 4, 5}, {0, 0, 0, 2, 2, 6}, {0, 0, 0, 2, 3, 5}, {0, 0, 0, 
   2, 4, 4}, {0, 0, 0, 3, 3, 4}, {0, 0, 1, 1, 1, 7}, {0, 0, 1, 1, 2, 
   6}, {0, 0, 1, 1, 3, 5}, {0, 0, 1, 1, 4, 4}, {0, 0, 1, 2, 2, 5}, {0,
    0, 1, 2, 3, 4}, {0, 0, 1, 3, 3, 3}, {0, 0, 2, 2, 2, 4}, {0, 0, 2, 
   2, 3, 3}, {0, 1, 1, 1, 1, 6}, {0, 1, 1, 1, 2, 5}, {0, 1, 1, 1, 3, 
   4}, {0, 1, 1, 2, 2, 4}, {0, 1, 1, 2, 3, 3}, {0, 1, 2, 2, 2, 3}, {0,
    2, 2, 2, 2, 2}, {1, 1, 1, 1, 1, 5}, {1, 1, 1, 1, 2, 4}, {1, 1, 1, 
   1, 3, 3}, {1, 1, 1, 2, 2, 3}, {1, 1, 2, 2, 2, 2}}}

{0.0093, {{0, 0, 0, 0, 0, 10}, {0, 0, 0, 0, 1, 9}, {0, 0, 0, 0, 2, 
   8}, {0, 0, 0, 0, 3, 7}, {0, 0, 0, 0, 4, 6}, {0, 0, 0, 0, 5, 5}, {0,
    0, 0, 1, 1, 8}, {0, 0, 0, 1, 2, 7}, {0, 0, 0, 1, 3, 6}, {0, 0, 0, 
   1, 4, 5}, {0, 0, 0, 2, 2, 6}, {0, 0, 0, 2, 3, 5}, {0, 0, 0, 2, 4, 
   4}, {0, 0, 0, 3, 3, 4}, {0, 0, 1, 1, 1, 7}, {0, 0, 1, 1, 2, 6}, {0,
    0, 1, 1, 3, 5}, {0, 0, 1, 1, 4, 4}, {0, 0, 1, 2, 2, 5}, {0, 0, 1, 
   2, 3, 4}, {0, 0, 1, 3, 3, 3}, {0, 0, 2, 2, 2, 4}, {0, 0, 2, 2, 3, 
   3}, {0, 1, 1, 1, 1, 6}, {0, 1, 1, 1, 2, 5}, {0, 1, 1, 1, 3, 4}, {0,
    1, 1, 2, 2, 4}, {0, 1, 1, 2, 3, 3}, {0, 1, 2, 2, 2, 3}, {0, 2, 2, 
   2, 2, 2}, {1, 1, 1, 1, 1, 5}, {1, 1, 1, 1, 2, 4}, {1, 1, 1, 1, 3, 
   3}, {1, 1, 1, 2, 2, 3}, {1, 1, 2, 2, 2, 2}}}

True

So the method based on IntegerPartitions is significantly faster. In order to get complete results:

f1[6, 10] // Map[Permutations] // Apply[Join]; // 
  RepeatedTiming
{0.00017, Null}
f1[4, 2] // Map[Permutations] // Apply[Join] //
  RepeatedTiming
{0.000012, {{0, 0, 0, 2}, {0, 0, 2, 0}, {0, 2, 0, 0}, {2, 0, 0, 
   0}, {0, 0, 1, 1}, {0, 1, 0, 1}, {0, 1, 1, 0}, {1, 0, 0, 1}, {1, 0, 
   1, 0}, {1, 1, 0, 0}}}
$\endgroup$
4
  • $\begingroup$ Both f1[2,4] and f2[2,4] give {{0, 4}, {1, 3}, {2, 2}} but the correct answer is {{2, 0, 0, 0}, {0, 2, 0, 0}, {0, 0, 2, 0}, {0, 0, 0, 2}, {1, 1, 0, 0}, {1, 0, 1, 0}, {1, 0, 0, 1}, {0, 1, 1, 0}, {0, 1, 0, 1}, {0, 0, 1, 1}} $\endgroup$ – geom Jan 19 at 19:33
  • $\begingroup$ @geom Oh, I put $N$ in front, so you should use f1[4, 2]. If you want to get complete results, try f1[4, 2] // Map[Permutations /* Reverse] // Flatten[#, 1] & :) $\endgroup$ – SneezeFor16Min Jan 19 at 19:42
  • $\begingroup$ I see, it would be helpful if you could modify your answer (for future readers :) ) $\endgroup$ – geom Jan 19 at 19:44
  • $\begingroup$ @geom You're right. I've updated! $\endgroup$ – SneezeFor16Min Jan 19 at 19:57

Your Answer

By clicking “Post Your Answer”, you agree to our terms of service, privacy policy and cookie policy

Not the answer you're looking for? Browse other questions tagged or ask your own question.