Sum of two elements of a list equaling a target

Question

In effort to learn OCaml, I'm going through some easy LeetCode problems and attempting to implement them. I arrive at the correct answer, but I'm still developing an intuition for canonical OCaml code. I'd appreciate any feedback on my implementation.

As mentioned, this is problem 1 on LeetCode. A quick summary of the problem: given an array of integers, return the indices of the 2 numbers that sum to a given "target".

For example, given nums: [2, 7, 11, 15] and target: 13, the result should be [0, 2] (the integers at indices 0 and 2 sum to 13).

Furthermore, we can assume a unique solution exists for the given array and target.

Here is my refined attempt at implementing this in OCaml:

let two_sum nums target =
  let rec aux i j =
    let n1 = List.nth nums i in
    let n2 = List.nth nums j in
      if n1 + n2 = target then [i; j]
      else if j = (List.length nums) - 1 then aux (i + 1) (i + 2)
      else aux i (j + 1)
  in aux 0 1;;

I appreciate any feedback.

Answer 1

Since the task is about using arrays, I would suggest to use OCaml arrays and not list. Because of lists, you need to use List.nth and that is quite costly.

Answer 2

You don't need to use arrays to do this efficiently.

let find_two_that_sum lst sum =
  let rec f lst sum =
    match lst with 
    | [] | [_] -> None
    | ((i, x)::xs) -> 
      match List.find (fun (_, x') -> x' = sum - x) xs with
      | (i', _) -> Some (i, i')
      | exception Not_found -> f xs sum
  in
  let lst' = List.mapi (fun i x -> (i, x)) lst in
  f lst' sum

This has type int list -> int -> (int * int) option

The inner function f deals with a list of tuples, where each tuple is both index and value.

If the list has zero or one element, there cannot be two elements that add up to the sum, so the result is None.

Otherwise we look at the first element, and look at the rest of the list to find a complement that adds up to the desired sum. If we find it, we can immediately return the two indices.

If we don't we try the same thing with the tail of the list. Obviously nothing plus the first number equals the sum, so we no longer need to consider that.

The recursion naturally ends when we either find two numbers (and their indices) that add up to the desired sum, or run out of numbers to consider.

We might also create a helper function that behaves like List.find but finds the index of a matching element in a list, or raises the Not_found exception.

let rec findi i f = function
  | [] -> raise Not_found
  | x::_ when f x -> i
  | _::xs -> findi (i+1) f xs

In many cases we would use an inner auxiliary function to remove the need to specify the starting index, but in this case, that will be useful as it removes the need to map indices to our list elements.

let find_two_that_sum lst sum =
  let rec f i lst sum =
    match lst with 
    | []  | [_] -> None
    | x::xs -> 
      try
        Some (i, findi (i+1) (fun x' -> x' = sum - x) xs)
      with Not_found -> 
        find_two_that_sum (i+1) xs sum 
  in  
  f 0 lst sum

Maps

Ideally we'd use a map for the best efficiency. As we iterate we build a map of values to indices. This is more efficient because we can insert into and lookup from maps in O(log n) time due to the fact that maps are built on a balanced binary search tree.

let two_sum target lst =
  let module Int_map = Map.Make (Int) in
  let rec aux map = function
    | [] -> None
    | (x, i)::xs ->
      (match Int_map.find_opt (target - x) map with
       | Some j -> Some (j, i)
       | None -> aux (Int_map.add x i map) xs)
  in
  aux Int_map.empty (List.mapi (fun i x -> (x, i)) lst)

But this still requires that we iterate over the entire list first to transform it into a new list with the elements and their indices. With sequences we can avoid having to do this.

let two_sum target lst =
  let module Int_map = Map.Make (Int) in
  let rec aux map seq =
    match seq () with
    | Seq.Nil -> None
    | Seq.Cons ((x, i), seq) ->
      (match Int_map.find_opt (target - x) map with
       | Some j -> Some (j, i)
       | None -> aux (Int_map.add x i map) seq)
  in
  lst
  |> List.to_seq
  |> Seq.mapi (fun i x -> (x, i))
  |> aux Int_map.empty