groupby.ml

open Core.Std


let maybe_list this that =
  match this with
  | Some l ->
     Some (List.append l [that]);
  | None -> Some [that]
  
let group_by f l =
  let ht = Hashtbl.Poly.create () in
  List.fold
    ~f:(fun () el ->
        let res = f el in
        let cur = Hashtbl.Poly.find ht res in
        Hashtbl.Poly.replace ht res (maybe_list cur el);
        ())
    ~init:() l;
  ht

open Core_kernel.Std

let group_by f l =
  let ht = Hashtbl.Poly.create () in
  List.fold
    l
    ~init:()
    ~f:(fun () el ->
        let res = f el in
        let group =
          match Hashtbl.Poly.find ht res with
          | Some els -> el :: els
          | None     -> [el]
        in
        Hashtbl.Poly.replace ht res group
    );
    ht