(* balanced.ml *)

(* definitions:

   a list xs is a PREFIX of a list ys iff xs @ us = ys for some us

   a list xs is a SUFFIX of a list ys iff us @ xs = ys for some us

   an int list ns is NICE iff ns has no nonempty sublists of sum 0 *)

(* val insert : 'a -> 'a list -> 'a list

   if ys is sorted in strictly ascending order and x appears in ys,
   then insert x ys = ys

   if ys is sorted in strictly ascending order and x does not appear
   in ys, then insert x ys is the result of putting x :: ys into
   strictly ascending order *)

let insert x =
      (* val ins : 'a list -> 'a list

         if ys is sorted in strictly ascending order and x appears in ys,
         then ins ys = ys

         if ys is sorted in strictly ascending order and x does not appear
         in ys, then ins ys is the result of putting x :: ys in
         strictly ascending order *)
      let rec ins ys =
            match ys with
                  []            -> [x]
                | y :: ys as zs ->
                    match compare x y with
                          -1 -> x :: zs
                        | 0  -> zs
                        | _  -> y :: ins ys
      in ins

(* val splitRev : int -> 'a list -> 'a list * 'a list

   if 0 <= n <= List.length xs, then splitRev n xs = (List.rev ys,
   zs), where ys is the prefix of xs of length n and zs is the suffix
   of xs of length (List.length xs - n) *)

let splitRev n xs =
      (* val split : int -> 'a list -> 'a list -> 'a list * 'a list

         if 0 <= n <= List.length xs, then split n ys xs = (List.rev
         us @ ys, vs), where us is the prefix of xs of length n and vs
         is the suffix of xs of length (List.length xs - n)

         termination: recursion is on xs *)

      let rec split n ys xs =
            if n = 0
            then (ys, xs)
            else split (n - 1) (List.hd xs :: ys) (List.tl xs)
      in split n [] xs

(* val scanZeroSum : int -> int list -> int option

   if there is a prefix of xs such that n + the prefix's sum is 0,
   then scanZeroSum n xs returns Some of the length of the
   shortest such prefix; otherwise, it returns None *)

let scanZeroSum n xs =
      (* val scan : int -> int -> int list -> int option

         if there is a prefix of xs such that n + the prefix's sum is
         0, then scan n i xs returns Some of i + the length of the
         shortest such prefix; otherwise it returns None

         termination: recursion is on xs *)

      let rec scan n i xs =
            if n = 0
            then Some i
            else match xs with
                       []      -> None
                     | x :: xs -> scan (n + x) (i + 1) xs
      in scan n 0 xs

(* val extend : int -> int list -> (int list)option * int list

   if ns : int list and the reversals of the prefixes of xs are the
   nice suffixes of ns, then extend n xs returns (opt, ys) where:

     if there is a balanced suffix of ns @ [n], then opt = Some ms,
     where ms is the unique such suffix; otherwise opt = None,

   and

     the reversals of the prefixes of ys are the nice suffixes of
     ns @ [n] *)

let extend n xs =
      match scanZeroSum n xs with
            None   -> (None, n :: xs)
          | Some i ->
              let (us, vs) = splitRev i (n :: xs)
              in (Some(List.hd vs :: us), List.rev us)

(* val balanced : 'a list -> 'a list list

   balanced ns returns the balanced sublists of ns, listed in
   strictly ascending order *)

let balanced ns =
      (* val bal : int list -> int list -> int list list -> int list list

         if ms is a suffix of ns, and the reversals of the prefixes of
         xs are the nice suffixes of the prefix of ns of length
         List.length ns - List.length ms, and zss is the balanced
         sublists of the prefix of ns of length List.length ns -
         List.length ms, listed in strictly ascending order, then bal
         ms xs zss is all the balanced sublists of ns, listed in
         strictly ascending order

         termination: recursion is on ms *)

      let rec bal ms xs zss =
            match ms with
                  []      -> zss
                | m :: ms ->
                    match extend m xs with
                          (None,    ys) -> bal ms ys zss
                        | (Some zs, ys) -> bal ms ys (insert zs zss)
      in bal ns [] []