Filtering an array in OCaml

Question

I have an array of integers where I'm treating 0 in a special way. I need to manipulate the array so that the output is an array of the same size with all the zeros after all the non-zero values. For example:

input:  [| 10; 5; 4; 0; 6; 0; 0; 0; 0 |]
output: [| 10; 5; 4; 6; 0; 0; 0; 0; 0 |]

This is some code I got to work:

let a = [| 10; 5; 4; 0; 6; 0; 0; 0; 0 |] in
let l = List.filter (fun c -> c <> 0) (Array.to_list a) in
let cells i = if i <= List.length l - 1 then List.nth l i else 0 in
Array.init (Array.length a) cells;;

I have a feeling that this solution could be improved by addressing:

1. converting back and forth between lists and arrays. Is there a better way than calling Array.to_list and then List.filter to keep the order and exclude certain values?
2. accessing a list using List.nth which is O(i) for each position in the array. Is there a better way to use a list's values as the first source when populating an array?

Answer 1

You're right that the List.nth in this code will cause it to be slow for large arrays.

You could construct the result as a list just by appending the right number of 0s to it. Then you could use Array.of_list to get the final array. That would avoid the quadratic cost of the repeated calls to List.nth.

You could also construct your result as an array directly without using lists at all.

Something like this:

let compress ar =
    let res = Array.make (Array.length ar) 0 in
    let c ix ari =
        if ari = 0 then ix
        else (res.(ix) <- ari; ix + 1)
    in
    let _ = Array.fold_left c 0 ar in
    res

This isn't pure code; i.e., it modifies the result array. That's kind of a drag. But arrays in OCaml are mutable after all.

Here is how it works on your test data:

# compress [| 10; 5; 4; 0; 6; 0; 0; 0; 0 |];;
- : int array = [|10; 5; 4; 6; 0; 0; 0; 0; 0|]