/* 

  
  3SUM (Three Elements That Sum To Zero) in Picat.
  
  From
  http://nathanleclaire.com/blog/2013/10/22/three-elements-that-sum-to-zero/
  """
  Given a collection of integers, return the indices of any three elements which sum to zero. 
  For instance, if you are given {-1, 6, 8, 9, 10, -100, 78, 0, 1}, you could return {0, 7, 8} 
  because -1 + 1 + 0 == 0. You can’t use the same index twice, and if there is no match you 
  should return {-1, -1, -1}.
  """

  Also see: https://en.wikipedia.org/wiki/3SUM

  This model solves the general (or rather general) N Sum problem.


  This Picat model was created by Hakan Kjellerstrand, hakank@gmail.com
  See also my Picat page: http://www.hakank.org/picat/

*/


import cp.


main => go.


go =>
   % Nums = [-1, 6, 8, 9, 10, -100, 78, 0, 1], % SAT
   Nums = [-10, -8, -6, 0, 3, 5, 18],
   N = 3, % The number of elements that should add to 0

   (n_sum(Nums, N, X, Ix) -> 
      println('nums'=Nums),
      println('   x'=X),
      println('  ix'=Ix)
     ;
      println([-1 : _I in 1..N])
   ),

   nl.

go1 =>
   Nums = [1, 6, 8, 9, 10, 100, 78, 0, 1], % UNSAT
   N = 3, % The number of elements that should add to 0

   (n_sum(Nums, N, X, Ix) -> 
      println('nums'=Nums),
      println('   x'=X),
      println('  ix'=Ix)
     ;
      println([-1 : _I in 1..N])
   ),

   nl.

go2 => 
   Len = 1000,
   Nums = random_list(Len,Len),
   println(nums=Nums),
   foreach(N in 2..Len)
      println(n=N),
      (n_sum(Nums, N, X, Ix) -> 
         % println('nums'=Nums),
         println('   x'=X),
         println('  ix'=Ix),
         nl
       ;
         println([-1 : _I in 1..N])
      )    
   end,
   nl.

%
% This is the general NSum
%
n_sum(Nums, N, X, Ix) => 

   Len = Nums.length,

   % decision variables
   X = new_list(Len), 
   X :: 0..1,

   % constraints
   sum([Nums[I]*X[I] : I in 1..Len]) #= 0,
   sum(X) #= N, 

   solve([], X),

   Ix = [I : I in 1..Len, X[I] = 1].



% Generates a list of both positive values
random_list(N, MaxVal) = [ (MaxVal div 2)-(random2() mod MaxVal) : _I in 1..N].