Rank Transform of a Matrix Leetcode in OCaml (https://leetcode.com/problems/rank-transform-of-a-matrix/)

rank.ml

type constr = Equal | Less | Greater 

let adjacents matrix (x,y) = 
  matrix 
  |> Array.mapi (fun ri row -> row |> Array.mapi (fun ci _col -> ri,ci))
  |> Array.to_list
  |> Array.concat
  |> Array.to_list 
  |> List.filter (fun (ri, ci) -> (ri = x) <> (ci = y))

let index matrix (x, y) = matrix.(x).(y) 

let rec propagate matrix rank node = 
  let n = index matrix node in 
  let nrank = index rank node in
  let update (x, y) value =
    rank.(x).(y) <- value ;
    propagate matrix rank (x,y)
  in 
  let f adj =
    let a = index matrix adj in
    let arank = index rank adj in
    let apply c  = 
      match c, compare nrank arank with
      | Less, -1 -> ()
      | Less, _ -> update adj (nrank+1)

      | Greater, 1 -> ()
      | Greater, _ -> update node (arank+1)

      | Equal, 0 -> ()
      | Equal, 1 -> update adj nrank 
      | Equal, -1 -> update node arank
      | Equal, _ -> assert false 
    in 
    match compare n a with 
    | -1 -> apply Less
    | 1 -> apply Greater
    | 0 -> apply Equal
    | _ -> assert false 
  in
  List.iter f (adjacents matrix node)

let matrix_rank_transform matrix =
  let rank = matrix |> Array.map @@ Array.map @@ Fun.const 1 in
  let nodes = 
    matrix 
    |> Array.mapi (fun j row -> row |> Array.mapi (fun i _ -> (j, i)) )
    |> Array.map Array.to_list
    |> Array.to_list
    |> List.concat 
  in 
  List.iter (propagate matrix rank) nodes ; 
  rank