If you don't really care about specific order given by Permutations
of IntegerPartitions
and can accept lexical order, then there's relatively easy algorithm calculating composition index.
Let's start with function generating all weak compositions in lexicographic order:
weakCompositionsLex // ClearAll
weakCompositionsLex[n_Integer, k_Integer?NonNegative] :=
Sort[Join @@ (Permutations /@ IntegerPartitions[n, {k}, Range[0, n]])]
Index of weak $k$-composition $c = (c_1, c_2, \ldots, c_k)$, in lexicographic order, can be calculated using following formula:
(formula = Binomial[Sum[Indexed[c, j], {j, 1, k}] + k - 1, k - 1]
- Sum[Binomial[Sum[Indexed[c, j], {j, i + 1, k}] - 1 + k - i, k-i], {i, 1, k - 1}]
) // TraditionalForm
$$
\binom{\sum_{j=1}^k c_j + k-1}{k-1}-\sum_{i=1}^{k-1} \binom{\sum_{j=i+1}^k c_j-1 + k-i}{k-i}
$$
Let's check that it gives correct results for all compositions with n
and k
up to 9
:
tmp = Evaluate@formula & /. {Indexed[c, i_] :> #[[i]]};
Table[tmp /@ weakCompositionsLex[n, k] === Range@Binomial[n + k - 1, k - 1], {n, 0, 9}, {k, 0, 9}];
And @@ Flatten@%
(* True *)
Above formula can be implemented as following top level function:
weakCompositionLexIndexTop = If[# === {},
1
(* else *),
With[{acc = Accumulate@Reverse@#, range = Range[Length@# - 1]},
Binomial[Last@acc + Length@# - 1, Length@# - 1] -
Total@Binomial[Most@acc + range - 1, range]
]
]&;
Basic tests:
Table[weakCompositionLexIndexTop /@ weakCompositionsLex[n, k] === Range@Binomial[n + k - 1, k - 1], {n, 0, 10}, {k, 0, 10}];
And @@ Flatten@%
(* True *)
We can also implement it as compiled function:
weakCompositionLexIndexC = Hold@Compile[{{comp, _Integer, 1}},
Module[{k, m, index, sum, tmpSum, jMin, jMax, bin, binPrev, l, i, iZeroSum, d, el},
k = Length@comp;
If[k <= 1, Return@1];
index = binPrev = comp[[k]];
sum = index - 1;
i = k - 1;
m = 1;
If[sum === -1,
While[comp[[i]] === 0 && i >= 2, --i];
m += k - 1 - i;
];
tmpSum = sum + comp[[i]];
If[tmpSum === 0 && i >= 2,
iZeroSum = i;
While[tmpSum === 0 && i >= 2, tmpSum += comp[[--i]]];
d = iZeroSum - i;
m += d;
index += d;
sum = 0;
binPrev = 1;
];
While[i >= 2,
++m;
el = comp[[i]];
tmpSum = sum;
sum += el;
jMin = 2;
If[m < sum,
jMax = m;
bin = l = sum + 1;
(* else *),
jMax = sum;
bin = l = m + 1;
];
If[el < jMax - jMin && binPrev > 0,
jMin = tmpSum + 1;
jMax = sum;
l = tmpSum + m;
bin = Quotient[binPrev l, m];
];
Do[bin = Quotient[bin (++l), j], {j, jMin, jMax}];
binPrev = bin;
index += bin;
--i;
];
sum += First@comp + 1;
If[sum === 0, Return@1];
jMax = If[m < sum, m, sum];
bin = l = sum + m;
Do[bin = Quotient[bin (--l), j], {j, 2, jMax}];
bin - index
],
RuntimeOptions -> {"Speed", "CatchMachineIntegerOverflow" -> True, "WarningMessages" -> False},
CompilationTarget -> "C", RuntimeAttributes -> {Listable}, Parallelization -> True
] /.
Part -> Compile`GetElement //. HoldPattern[Compile`GetElement@x__ = y_] :> (Part@x = y) //.
HoldPattern[Plus][pre___, x_, HoldPattern[Times][-1, y_], post___] :> Plus[pre, Subtract[x, y], post] //
ReleaseHold;
Test all compositions with n
and k
up to 12
:
Table[weakCompositionLexIndexC@weakCompositionsLex[n, k] === Range@Binomial[n + k - 1, k - 1], {n, 0, 12}, {k, 12}];
And @@ Flatten@%
(* True *)
Final function that uses compiled function when possible, and falls back to top level version on integer overflow and for symbolic arguments.
With[{cf = ReplacePart[weakCompositionLexIndexC, 7 -> weakCompositionLexIndexTop]},
weakCompositionLexIndex // ClearAll;
weakCompositionLexIndex@l_List := cf@l;
]
Tests:
weakCompositionLexIndex@{}
(* 1 *)
weakCompositionLexIndex@{a, b, c, d}
(* 1 + c - 1/2 (1 + c + d) (2 + c + d) + 1/2 (1 + b + c + d) (2 + b + c + d) - 1/6 (1 + b + c + d) (2 + b + c + d) (3 + b + c + d) + 1/6 (1 + a + b + c + d) (2 + a + b + c + d) (3 + a + b + c + d) *)
% /. {a -> 2, b -> 0, c -> 1, d -> 5}
(* 83 *)
weakCompositionLexIndex@{2, 0, 1, 5}
(* 83 *)
weakCompositionLexIndex@Join[{1}, ConstantArray[0, 10^6]]
(* 1000001 *)
weakCompositionLexIndex@{210, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0}
(* 2620460268444674457614627940 *)
weakCompositionLexIndex@{64, 6, 0, 8*10^20 - 64 - 6}
(* 20479999999999999998468800000000000000039320 *)
weakCompositionLexIndex@c
(* weakCompositionLexIndex[c] *)
Table[weakCompositionLexIndex@weakCompositionsLex[n, k] === Range@Binomial[n + k - 1, k - 1], {n, 0, 12}, {k, 12}];
And @@ Flatten@%
(* True *)
Let's check time and memory usage on data analogous to one used in answer by Henrik Schumacher.
compositions = weakCompositionsLex[400, 4];
rCompositions = RandomChoice[compositions, 100000];
weakCompositionLexIndex@rCompositions; // MaxMemoryUsed // RepeatedTiming
(* {0.021, 801872} *)
We can index all ten millions compositions and check that we get correct consecutive numbers:
res = weakCompositionLexIndex@compositions; // MaxMemoryUsed // RepeatedTiming
res === Range@Binomial[400 + 4 - 1, 4 - 1]
(* {2.36, 86621048} *)
(* True *)