Combinations of size k from a list in OCaml

Question

I wanted to write a small but non-trivial function in OCaml to test my understanding. I have only a basic understanding of the libraries, and no idea at all of what is considered good style. The function computes all combinations of size k from a list, where 0 <= k <= length lst.

Would someone mind commenting on

1. Whether the function is generally well written (have I covered all cases, are there opportunities for tail recursion that I've missed?)

2. Have I made good use of libraries (e.g. I defined is_empty and tails because I couldn't find them in the List module, but maybe they are somewhere else?)

3. Is the style okay, particularly the use of let statements and indentation?

The code is:

let rec tails = function
  | []          -> []
  | _ :: t as l -> l :: tails t

let is_empty = function
  | [] -> true
  | _  -> false

let rec combnk k lst =
  if k = 0 then [[]]
  else let f = function
    | []      -> [] (* I think this is unnecessary, but I get a pattern match warning o/w *)
    | x :: xs -> List.map (fun z -> x :: z) (combnk (k-1) xs)
  in if is_empty lst then []
     else List.concat (List.map f (tails lst))

---

Based on the excellent comments by amon (see below) I have written to what I think is the most readable version of this function, which is the one that inlines the definition of tails and gets rid of is_empty completely, but doesn't go all the way to removing the use of List.concat and List.map, because I believe in using library functions to simplify the code wherever possible.

In particular, the layout guidelines make the structure of the function much clearer, and I think that in this version it is obvious what algorithm is being used, whereas it is somewhat obfuscated in the original. Thanks, amos!

let rec combnk k lst =
  if k = 0 then
    [[]]
  else
    let rec inner = function
      | []      -> []
      | x :: xs -> List.map (fun z -> x :: z) (combnk (k - 1) xs) :: inner xs in
    List.concat (inner lst)

Answer 1

Your tails is very beatiful, your is_empty useless, and combnk a mess.

In combnk, your indentation obfuscates the actual structure of the code. Here is a better indentation:

let rec combnk k lst =
  if k = 0 then
    [[]]
  else
    let f = function
      | []      -> []
      | x :: xs -> List.map (fun z -> x :: z) (combnk (k - 1) xs)
    in
      if is_empty lst then
        []
      else
        List.concat (List.map f (tails lst))

Now there are some interesting observations to be made here: The branch if is_empty lst then [] does not access f, so we could move this test outside of the let:

let rec combnk k lst =
  if k = 0 then
    [[]]
  else if is_empty lst then
    []
  else
    let f = function
      | []      -> []
      | x :: xs -> List.map (fun z -> x :: z) (combnk (k - 1) xs)
    in
      List.concat (List.map f (tails lst))

But is the is_empty test actually necessary? tails [] produces an empty list, List.map f [] produces an empty list for any function f, and List.concat [] also produces an empty list. We now have:

let rec combnk k lst =
  if k = 0 then
    [[]]
  else
    let f = function
      | []      -> []
      | x :: xs -> List.map (fun z -> x :: z) (combnk (k - 1) xs)
    in
      List.concat (List.map f (tails lst))

How can this be improved? We can move the tails definitions inside the else-branch let so that it's restricted to the only scope where it is used:

let rec combnk k lst =
  if k = 0 then
    [[]]
  else
    let rec tails = function
      | []          -> []
      | _ :: t as l -> l :: tails t
    and f = function
      | []      -> []
      | x :: xs -> List.map (fun z -> x :: z) (combnk (k - 1) xs)
    in
      List.concat (List.map f (tails lst))

Regarding the question whether the case [] -> [] in f is necessary except for the type system: The answer is no, as the list produced by tails cannot contain another empty list – l :: [] is l again. This would change when you swap [] -> [] in tails for [] -> [[]], which would be arguably more correct.

Now that f and tails are so close together you may notice some similarities. Indeed, we can combine the two directly, thus getting rid of one map:

let rec combnk k lst =
  if k = 0 then
    [[]]
  else
    let rec inner = function
      | []      -> []
      | x :: xs -> (List.map (fun z -> x :: z) (combnk (k - 1) xs)) :: inner xs
    in
      List.concat (inner lst)

Of course, inner could be made partially tail recursive (but this reverses the order of combinations):

let rec combnk k lst =
  if k = 0 then
    [[]]
  else
    let rec inner acc = function
      | []      -> acc
      | x :: xs ->
        let this_length = List.map (fun z -> x :: z) (combnk (k - 1) xs)
        in
          inner (this_length :: acc) xs
    in
      List.concat (inner [] lst)

There is still an indirect recursion through combnk, more obvious if we rewrite it like this:

let rec combnk k lst =
  let rec inner acc k lst =
    match k with
    | 0 -> [[]]
    | _ ->
      match lst with
      | []      -> List.flatten acc
      | x :: xs ->
        let this_length = List.map (fun z -> x :: z) (inner [] (k - 1) xs)
        in
          inner (this_length :: acc) k xs
  in
    inner [] k lst

Now all that is left to do is to write a map that takes an external accumulator, thus also removing the need for flatten or concat:

let rec combnk k lst =
  let rec inner acc k lst =
    match k with
    | 0 -> [[]]
    | _ ->
      match lst with
      | []      -> acc
      | x :: xs ->
        let rec accmap acc f = function
          | []      -> acc
          | x :: xs -> accmap ((f x) :: acc) f xs
        in
          let newacc = accmap acc (fun z -> x :: z) (inner [] (k - 1) xs)
          in
            inner newacc k xs
    in
      inner [] k lst

Answer 2

It seems highly idiomatic to view this as an exercise in pattern-matching. There are fundamentally four scenarios to handle which we can model by pattern-matching on both k and lst in one match expression.

• We're looking for k element combinations of an empty list. Obviously there aren't any.
• We're looking for k element combinations of any list where k is 0. There also aren't any in this scenario.
• We're looking for k element combinations of any list, where k is 1. We simply need to map each element of the list to a singleton list containing that element.
• Otherwise we're looking for k (where k is greater than 1) element combinations of a non-empty list. This should look familiar.

let rec combnk k lst =
  match k, lst with
  | _, [] | 0, _ -> []
  | 1, _ -> List.map (fun x -> [x]) lst
  | _, x::xs -> 
    let k' = k - 1 in
    List.map (fun tl -> x :: tl) (combnk k' xs) @ combnk k xs

Though the above will return an empty list for negative values of k, you may wish to explicitly raise an exception if a negative value is provided for k.

let rec combnk k lst =
  if k < 0 then invalid_arg "combnk: negative length";
  match k, lst with
  | _, [] | 0, _ -> []
  | 1, _ -> List.map (fun x -> [x]) lst
  | _, x::xs -> 
    let k' = k - 1 in
    List.map (fun tl -> x :: tl) (combnk k' xs) @ combnk k xs

Repeats

While this is a fairly straightforward algorithm for finding combinations, it features numerous calls with the same arguments. This suggests that there is an opportunity for memoization.

If we modify the function with a print, we can see these repeated calls.

let rec combnk k lst =
  if k < 0 then invalid_arg "combnk: negative length";
  let pp_print_sep fmt () = Format.fprintf fmt "; " in
  let pp_print_int_list = Format.(
    pp_print_list pp_print_int ~pp_sep: pp_print_sep
  ) in
  Format.printf "%d [%a]\n" k pp_print_int_list lst;
  match k, lst with
  | _, [] | 0, _ -> []
  | 1, _ -> List.map (fun x -> [x]) lst
  | _, x::xs -> 
    let k' = k - 1 in
    let sublst1 = List.map (fun tl -> x :: tl) (combnk k' xs) in
    let sublst2 = combnk k xs in
    sublst1 @ sublst2

# combnk 3 [1; 2; 3; 4; 5; 6];;
3 [1; 2; 3; 4; 5; 6]
2 [2; 3; 4; 5; 6]
1 [3; 4; 5; 6]
2 [3; 4; 5; 6]
1 [4; 5; 6]
2 [4; 5; 6]
1 [5; 6]
2 [5; 6]
1 [6]
2 [6]
1 []
2 []
3 [2; 3; 4; 5; 6]
2 [3; 4; 5; 6]
1 [4; 5; 6]
2 [4; 5; 6]
1 [5; 6]
2 [5; 6]
1 [6]
2 [6]
1 []
2 []
3 [3; 4; 5; 6]
2 [4; 5; 6]
1 [5; 6]
2 [5; 6]
1 [6]
2 [6]
1 []
2 []
3 [4; 5; 6]
2 [5; 6]
1 [6]
2 [6]
1 []
2 []
3 [5; 6]
2 [6]
1 []
2 []
3 [6]
2 []
3 []
- : int list list =
[[1; 2; 3]; [1; 2; 4]; [1; 2; 5]; [1; 2; 6]; [1; 3; 4]; [1; 3; 5]; [1; 3; 6];
 [1; 4; 5]; [1; 4; 6]; [1; 5; 6]; [2; 3; 4]; [2; 3; 5]; [2; 3; 6]; [2; 4; 5];
 [2; 4; 6]; [2; 5; 6]; [3; 4; 5]; [3; 4; 6]; [3; 5; 6]; [4; 5; 6]]

We can start memoizing by creating a map module to handle the arguments we're sending to combnk.

module Int_and_int_list_map = Map.Make (struct
  type t = int * int list
  let compare = compare
end)

We now need to pass a map into and out of this function via function arguments and return value and update it as we proceed to avoid unnecessary work.

let rec memo_combnk k lst map =
  if k < 0 then invalid_arg "combnk: negative length";

  let pp_print_sep fmt () = Format.fprintf fmt "; " in
  let pp_print_int_list = Format.(
    pp_print_list pp_print_int ~pp_sep: pp_print_sep
  ) in
  Format.printf "%d [%a]\n" k pp_print_int_list lst;

  match Int_and_int_list_map.find_opt (k, lst) map with
  | Some lst -> (lst, map)
  | None -> (
      match k, lst with
      | _, [] | 0, _ -> ([], map)
      | 1, _ -> 
        let results = List.map (fun x -> [x]) lst in
        (results, map |> Int_and_int_list_map.add (k, lst) results)
      | _, x::xs -> 
        let k' = k - 1 in
        let (tails, map) = memo_combnk k' xs map in
        let sublst1 = List.map (fun tl -> x :: tl) tails in
        let (sublst2, map) = memo_combnk k xs map in
        let full_list = sublst1 @ sublst2 in
        (full_list, map |> Int_and_int_list_map.add (k, lst) full_list)
    )

Of course, having to pass that empty map in and have it be part of the output is ugly, so we can wrap this in an outer function.

let combnk k lst =
  let rec memo_combnk k lst map =
    if k < 0 then invalid_arg "combnk: negative length";

    let pp_print_sep fmt () = Format.fprintf fmt "; " in
    let pp_print_int_list = Format.(
      pp_print_list pp_print_int ~pp_sep: pp_print_sep
    ) in
    Format.printf "%d [%a]\n" k pp_print_int_list lst;

    match Int_and_int_list_map.find_opt (k, lst) map with
    | Some lst -> (lst, map)
    | None -> (
        match k, lst with
        | _, [] | 0, _ -> ([], map)
        | 1, _ -> 
          let results = List.map (fun x -> [x]) lst in
          (results, map |> Int_and_int_list_map.add (k, lst) results)
        | _, x::xs -> 
          let k' = k - 1 in
          let (tails, map) = memo_combnk k' xs map in
          let sublst1 = List.map (fun tl -> x :: tl) tails in
          let (sublst2, map) = memo_combnk k xs map in
          let full_list = sublst1 @ sublst2 in
          (full_list, map |> Int_and_int_list_map.add (k, lst) full_list)
      )
  in
  memo_combnk k lst Int_and_int_list_map.empty
  |> fst

Now testing we can see the effect:

# combnk 3 [1; 2; 3; 4; 5; 6];;
3 [1; 2; 3; 4; 5; 6]
2 [2; 3; 4; 5; 6]
1 [3; 4; 5; 6]
2 [3; 4; 5; 6]
1 [4; 5; 6]
2 [4; 5; 6]
1 [5; 6]
2 [5; 6]
1 [6]
2 [6]
1 []
2 []
3 [2; 3; 4; 5; 6]
2 [3; 4; 5; 6]
3 [3; 4; 5; 6]
2 [4; 5; 6]
3 [4; 5; 6]
2 [5; 6]
3 [5; 6]
2 [6]
3 [6]
2 []
3 []
- : int list list =
[[1; 2; 3]; [1; 2; 4]; [1; 2; 5]; [1; 2; 6]; [1; 3; 4]; [1; 3; 5]; [1; 3; 6];
 [1; 4; 5]; [1; 4; 6]; [1; 5; 6]; [2; 3; 4]; [2; 3; 5]; [2; 3; 6]; [2; 4; 5];
 [2; 4; 6]; [2; 5; 6]; [3; 4; 5]; [3; 4; 6]; [3; 5; 6]; [4; 5; 6]]

Scalability

Of course, actually comparing two lists in the tuples in our map is an O(n) operation, and has a distinct cost, but so does the recursion which would otherwise be needlessly repeated. The larger the data set, the more improvement we'll see from memoizing.

k	n	combnk	memo_combnk	%		k	n	combk	memo_combnk	%
3	5	31	19	61.3%		4	5	51	25	49.0%
3	6	43	23	53.5%		4	6	83	31	37.3%
3	7	57	27	47.4%		4	7	127	37	29.1%
3	8	73	31	42.5%		4	8	185	43	23.2%
3	20	421	79	18.8%		4	20	2,701	115	4.26%
3	21	463	83	17.9%		4	21	3,123	121	3.87%
3	100	10,101	399	3.95%		4	100	333,501	595	0.178%
3	101	10,303	403	3.91%		4	101	343,603	601	0.175%
3	500	250,501	1,999	0.798%

Flexibility

That's all well and good, but it can only work on lists of int values because of Int_and_int_list_map. We can use first-class modules to make this more flexible.

let combnk (type a) (module T : Map.OrderedType with type t = a) k lst =
  let module M = Map.Make (struct
    type t = int * T.t list
    let compare = compare
  end) in  

  let rec memo_combnk k lst map =
    match M.find_opt (k, lst) map with
    | Some lst -> (lst, map)
    | None -> (
        match k, lst with
        | _, [] | 0, _ -> ([], map)
        | 1, _ -> 
          let results = List.map (fun x -> [x]) lst in
          (results, map |> M.add (k, lst) results)
        | _, x::xs -> 
          let k' = k - 1 in
          let (tails, map) = memo_combnk k' xs map in
          let sublst1 = List.map (fun tl -> x :: tl) tails in
          let (sublst2, map) = memo_combnk k xs map in
          let full_list = sublst1 @ sublst2 in
          (full_list, map |> M.add (k, lst) full_list)
      )
  in

  memo_combnk k lst M.empty
  |> fst

Testing this:

# combnk (module Char) 3 ['1'; '2'; '3'; '4'; '5'; '6'];;
- : char list list =
[['1'; '2'; '3']; ['1'; '2'; '4']; ['1'; '2'; '5']; ['1'; '2'; '6'];
 ['1'; '3'; '4']; ['1'; '3'; '5']; ['1'; '3'; '6']; ['1'; '4'; '5'];
 ['1'; '4'; '6']; ['1'; '5'; '6']; ['2'; '3'; '4']; ['2'; '3'; '5'];
 ['2'; '3'; '6']; ['2'; '4'; '5']; ['2'; '4'; '6']; ['2'; '5'; '6'];
 ['3'; '4'; '5']; ['3'; '4'; '6']; ['3'; '5'; '6']; ['4'; '5'; '6']]

Optimization

We've reduced the number of calls to memo_combnk but those calls are expensive now because of looking up in a map where part of the key is a list. In order to insert or lookup in a map, those lists have to be compared, which requires iterating over them.

We can optimize this by realizing that the length of the list can be calculated once and is then known since lists are immutable. Rather than using the k and the list we're currently searching as a key, we'll use the k and the offset into the list. A tuple of two integers is much cheaper to compare.

As a side benefit, this no longer requires first-class modules to be generally applicable, since our map doesn't need to know about the kind of data we're finding combinations of.

let combnk k lst =
  let module Int_pair = struct
    type t = int * int  
  let compare = compare
  end in
  let module M = Map.Make (Int_pair) in  
  
  let len = List.length lst in

  let rec memo_combnk k lst map offset =
    match M.find_opt (k, offset) map with
    | Some lst -> (lst, map)
    | None -> 
      if k = 0 || len - offset < k then 
        ([], map)
      else if k = 1 then 
        let results = lst |> List.map (fun x -> [x]) in
        let map = map |> M.add (k, offset) results in
        (results, map)
      else
        let k' = k - 1 in
        let offset' = offset + 1 in
        (* We know that lst will be non-empty because we've
         * already eliminated any situations where it's empty. 
         * Now we want to silence the warning about 
         * non-exhaustive pattern matching. 
         *)
        let[@warning "-8"] (x::xs) = lst in
        let (tails, map) = memo_combnk k' xs map offset' in
        let sublst1 = tails |> List.map (fun tl -> x :: tl) in
        let (sublst2, map) = memo_combnk k xs map offset' in
        let full_list = sublst1 @ sublst2 in
        let map = map |> M.add (k, offset) full_list in
        (full_list, map)
  in

  memo_combnk k lst M.empty 0 
  |> fst
```