let flatten_tail l =
let rec flat acc = function
| [] -> List.rev acc
| hd::tl -> flat (List.rev_append hd acc) tl
in
flat [] l
let concat_map_tail f l =
List.rev_map f l |> List.rev |> flatten_tail
let rm_dup l =
if List.length l = 0 then l
else
let sl = List.sort compare l in
List.fold_left (
fun (acc, e) x -> if x <> e then x::acc, x else acc,e
) ([List.hd sl], List.hd sl) (List.tl sl) |> fst |> List.rev
(* algorithm starts from here *)
let buckets m =
let rec generate acc m =
if m = 0 then acc
else generate (0::acc) (m-1)
in
generate [] m
let throw_1_ball bs =
let rec throw acc before = function
| [] -> acc
| b::tl ->
let new_before = b::before in
let new_acc = (List.rev_append before ((b+1)::tl))::acc in
throw new_acc new_before tl
in
throw [] [] bs
let throw_n_ball n m =
let bs = buckets m in
let rec throw i acc =
if i = 0 then acc
else throw (i-1) (concat_map_tail throw_1_ball acc |> rm_dup)
in
throw n [bs]
Above is the correct code, it is scary because I added several utility functions and make things as tail-recursive as possible. But the idea is very simple.
Here is the algorithm:
- Let's say we have 3 buckets, initially it is [0;0;0].
- If we throw 1 ball into the 3 buckets, we have 3 cases each of which
is a snapshot of the buckets, i.e., [[1;0;0];[0;1;0];[0;0;1]].
- Then if we have 1 more ball, for each case above, we will 3 cases,
so the resulting case list have 9 cases
- Then if we have 1 more ball, .....
In this way, we will generate 3^n
cases and many of them may be redundant.
So when generated each case list, we just remove all duplicates in the case list.
utop # throw_n_ball 3 2;;
- : int list list = [[0; 3]; [1; 2]; [2; 1]; [3; 0]]
utop # throw_n_ball 5 3;;
- : int list list = [[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]]