Speed up random set partition code

Question

I have managed to construct code for randomly generating a set partition of size n with k parts. Note that the total number of such objects is counted by a Stirling number of the second kind.

The code is very simple; there are two cases for the largest entry n. Either it is in its own block, or not. It is easy to compute the probability for the first case, and then one uses recursion.

RandomSetPartition[0, k_Integer] := Which[k == 0, {{}}, {}];
RandomSetPartition[1, k_Integer] := Which[k == 1, {{1}}, {}];
RandomSetPartition[n_Integer, 1] := {Range@n};
RandomSetPartition[n_Integer, n_Integer] := List /@ (Range@n);

RandomSetPartition[n_Integer, k_Integer] := If[
   RandomReal[] <= StirlingS2[n - 1, k - 1]/StirlingS2[n, k],
   Append[RandomSetPartition[n - 1, k - 1], {n}], 
   MapAt[Append[#, n] &, RandomSetPartition[n - 1, k], 
    RandomInteger[{1, k}]]
   ];

The issue is, the code is rather slow for large arguments, say RandomSetPartition[1000,30] takes about 2 seconds on my machine. I am not sure what the issue with my code is; the complexity should be O(n), but some tests indicate the complexity is higher for my code. What can I change to make it more efficient?

Answer 1

I rewrote this answer completely. The previous answer is available here.

In OP's code, the StirlingS2 calls are expensive for large $n$. For example, calling

StirlingS2[4500,2111]

for the first time takes about 1 second on my machine. From OPs comment I understand that they want to call RandomSetPartition with say $n=20000$ and with many different values of $k$. The solution presented here was designed with that case in mind.

Code. Run the following once, with nmax equal to the largest value that one is planning to use:

nmax=2000;
s2prev={1.};
Do[
   s2=Append[Range[n-1]*s2prev+Prepend[Most[s2prev],0.],1.];
   sq[n]=Quiet[Developer`ToPackedArray[N[s2prev/Rest[s2]]]];
   s2prev=s2;
,{n,2,nmax}];

(* this is equal to N[StirlingS2[n-1,k-1]/StirlingS2[n,k]] *)
squot[n_,1]:=0.;
squot[n_,k_]:=sq[n][[k-1]];

This takes about $2$ seconds for nmax=2000, and it takes about $6$ minutes for nmax=20000. Since this has to be done only once, it seems acceptable.

Random set partitions can then be generated using

rSow[n_Integer,n_Integer]:=Scan[Sow[#,#]&,Range[n]];
rSow[n_Integer,k_Integer]:=rSow[n-1,If[RandomReal[]<=squot[n,k],
  Sow[n,k];k-1,
  Sow[n,RandomInteger[{1,k}]];k]];
RandomSetPartition[n_Integer,k_Integer]:=Block[{$IterationLimit=Infinity},
  Flatten[Last[Reap[rSow[n,k],Range[k]]],1]];

Example:

SeedRandom[1];
RandomSetPartition[10,5]
(* {{10,8,1},{2},{9,6,5,3},{4},{7}} *)

Timing:

First[RepeatedTiming[RandomSetPartition[2000,1000]]]
(* 0.00909227 *)

First[RepeatedTiming[RandomSetPartition[20000,10000]]]
(* 0.111011 *)

Comments:

• The code for squot uses the recurrence relation for StirlingS2. It is also written to keep memory footprint reasonably low.
• The code for RandomSetPartition is a direct adaptation of OP's code. It uses Reap-Sow instead of Append. It also uses tail recursion.