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