(* 
   Steve The Gambler: Caught by the IRS
   A subset-sum problem with a constraint on the size of the subset to find
   cameron {at} theaboutbox {dot} com

   If you want to use this code, you can, without warranty.
*)

(* Steve's daily winnings *)
let daily_winnings = [-50; -21; 13; 171; 14; -42; -58; 109; 4; 7; -23; 
		      -44; -98; -121; 101;  33; 87; -121; -40; -65; 43; 
		      54; -45; -12; -12; 38; 25; 3; 7; 8];;

(* Splits a list into positive and negative sublists. Sorts the lists so the
   numbers closest to zero are first. This way, if there is no ideal list, we
   create the smallest non-ideal subset we can. *)
let rec split_list l pos neg = match l with
    [] -> ((List.sort compare pos),(List.rev (List.sort compare neg))) 
  | h::t when h < 0 -> split_list t pos (h::neg)
  | h::t -> split_list t (h::pos) neg;;

(* Sums a list *)      
let sum = List.fold_left (+) 0;;

(* Prints an int list. Thanks, Andrew! *)
let rec print_int_list = function
  | [] -> ()
  | x::xs -> print_int x; print_string " "; print_int_list xs;;

(*
  Subset sum solver that looks for a sublist of the given size that sums to zero.

  Splits the list into positive and negative sublists to limit the number of
  decisions at any point in searching the problem space. When the running sum is
  positive, choose negative numbers and when the running sum is negative, choose 
  positive numbers.

  We employ some tricks to limit the effort spent searching because we know the size
  of the subset we need to find. If we are deep enough in the search space that there 
  is only one outcome, we return it immediately. If we run out of positive numbers
  and the running sum is negative, we fill out the list with the smallest positive
  numbers we can and vice versa.
*)
let subset_sum list size =
  let (pos,neg) = split_list list [] [] in
  (* Returns true if this list is the right size and sums to zero *)
  let suitable l = (List.length l) = size && (sum l) = 0 in
  (* Returns the best list: the one whose sum is closest to zero *)
  let best_list l1 l2 = match (l1, l2) with
      ([], l2) -> l2
    | (l1, []) -> l1
    | (l1, l2) -> let s1 = abs (sum l1) and s2 = abs (sum l2) in
		    if s1 < s2 then l1 else l2 in
  let rec subset_aux l s p n = match (l,s,p,n) with
      (list, sum, _, _) when (List.length list) = size -> list
    (* Reached a dead end in the search space - only one possible outcome *)
    | (list, _, p, n) when (List.length list) + (List.length p) + (List.length n) = size ->
	list @ p @ n
    (* Have a positive sum and only negative numbers: pile on the smallest neg #s possible *)
    | (list, sum, [], n::t) when sum < 0 -> subset_aux (n::list) (sum+n) [] t 
    (* Have a negative sum and only positive numbers: pile on the smallest pos #s possible *)
    | (list, sum, p::t, []) when sum > 0 -> subset_aux (p::list) (sum+p) t []
    (* Negative sum: pick a positive number and recurse *) 
    | (list, sum, p::t, n) when sum <= 0 -> 
	let with_list = subset_aux (p::list) (sum+p) t n in
	  if (suitable with_list) then with_list
	  else let without_list = subset_aux list sum t n in
		 best_list with_list without_list
    (* Positive sum: pick a negative number and recurse *)
    | (list, sum, p, n::t) ->
	let with_list = subset_aux (n::list) (sum+n) p t in
	  if (suitable with_list) then with_list
	  else let without_list = subset_aux list sum p t in
		 best_list with_list without_list 
  in
    subset_aux [] 0 pos neg;;

let read_days =
  if (Array.length Sys.argv) = 1 then 9
  else int_of_string Sys.argv.(1);;

let _ = 
  let best_list = subset_sum daily_winnings read_days in
    print_string "Best List: ";
    print_int_list best_list;
    print_string "\n";;