Picking five numbers that sum to S

Question

Given an array A of N nonnegative numbers, I\'m interested in finding the number of ways you can pick 5 numbers (from distinct positions in the array)

Answer 1

I think the fact that the numbers must have distinct positions is a red herring. You can use the inclusion-exclusion principle to count the number of all positions (i,j,k,l,m) where x[i]+x[j]+x[k]+x[l]+x[m]=S and i,j,k,l,m are distinct:

 sums with i!=j,i!=k,i!=l...,l!=m  = all sums 
                                   - sums with i=j
                                   - ...
                                   - sums with l=m
                                   + sums with i=j and j=k
                                   + ...
                                   + sums with k=l and l=m
                                   - ...
                                   + sums with i=j=k=l=m

Computing the sums on the right, except the first one, is doable in O(N^2 log N). For example, to find the number of positions (i,i,k,l,m) such that x[i]+x[i]+x[k]+x[l]+x[m]=S you can create sorted arrays with sums {2a+b} and {c+d} and check if they have elements x,y such that x+y=S.

Main algorithm

So it's enough to compute how many are there positions (i,j,k,l,m) where x[i]+x[j]+x[k]+x[l]+x[m]=S and i,j,k,l,m are not necessarily different. Basically, you can use Moron's solution this way:

• Create a sorted array of sums {a+b: a,b are numbers from array}; group equal elements into one, remembering count. For example, for array [1,1,3] you get nine sums [2,2,2,2,4,4,4,4,6] of the form a+b. Then you group same elements remembering counts: [(2,4),(4,4),(6,1)]. This step is O(N^2 log N).

• For each e, count how many are there pairs of elements in the array that sum to S-e. As in Moron's solution, you have two pointers, one going right, one going left. If the sum is too low, move the first pointer increasing the sum; if the sum is too high, move the second pointer decreasing it.

  Suppose the sum is correct. This means one points to (a,x) and second to (b,y) where a+b=S-e. Increase the counter by x*y and move both pointers (You could move only one pointer, but on the next step there would be no match, and the second pointer would be moved then.).

For example, for [(2,4),(4,4),(6,1)] array and S-e=8, the first pointer points at (2,4) and the second at (6,1). Since 2+6=8, you add 4 and move both pointers. Now they both point at (4,4), so you increase the counter by 16. Don't stop! The pointers pass each other, and you get first at (6,1), second at (2,4), increase the counter by 4.

So, in the end, there are 4+16+4=24 ways to get 8 as a sum of 4 elements of [1,1,3]:

>>> len([k for k in itertools.product([1,1,3],repeat=4) if sum(k) == 8])
24

Prelude Control.Monad> length [k | k <- replicateM 4 [1,1,3], sum k == 8]
24

Repeating that for each e, you'll get count of ways to get S as a sum of 5 elements.

For [1,1,1,1,1] and S-e=4, the sums array would be [(2,25)], and you'd get that there are 625 ways to get sum of 4.

For each e, this step is linear in size of the array (so it's O(N²)), so the loop takes O(N³).

On inclusion-exclusion:

Call a quintuple (i,j,k,l,m) "proper" if x[i]+x[j]+x[k]+x[l]+x[m]=S. The goal is to count the number of proper quintuples (i,j,k,l,m) where i,j,k,l,m are pairwise distinct. The main algorithm can count in O(N^3) how many are there proper quintuples which have not necessarily distinct components. The remaining thing is to count those "wrong" tuples.

Consider the subsets of proper quintuples

A_xy={(i,j,k,l,m): indices on x-th and y-th place are the same}

For example, A₂₄ is the set of proper quintuples (i,j,k,l,m) where j=l.

The set of wrong quintuples is:

A₁₂ ∪ A₁₃ ∪ ... ∪ A₄₅

Counting its cardinality by inclusion-exclusion:

|A₁₂ ∪ A₁₃ ∪ ... ∪ A₄₅| = |A₁₂| + |A₁₃| + ... + |A₄₅| - |A₁₂ ∩ A₂₃| - ... - |A₃₄ ∩ A₄₅| + ... + |A₁₂ ∩ A₂₃ ∩ ... ∩ A₃₅ ∩ A₄₅|

There are 2¹⁰=1024 summands here. But a lot of the cardinalities is the same.

The only things you have to count is:

• X₁ = |A₁₂| - quintuples with i=j
• X₂ = |A₁₂ ∩ A₂₃| - quintuples with i=j=k
• X₃ = |A₁₂ ∩ A₂₃ ∩ A₃₄| - quintuples with i=j=k=l
• X₄ = |A₁₂ ∩ A₂₃ ∩ A₃₄ ∩ A₄₅| - quintuples with i=j=k=l=m
• X₅ = |A₁₂ ∩ A₃₄| - quintuples with i=j,k=l
• X₆ = |A₁₂ ∩ A₂₃ ∩ A₄₅| - quintuples with i=j=k,l=m

You can observe, by permuting, all other sets are represented here. For example, A₂₄ has the same cardinality as A₁₂.

Counting cardinalities of those 6 sets is rather easy. For the first one, you create arrays {2a+b} and {c+d} and count how many are there common elements; for the other ones, there are only 3 or less free variables, so even a simple loop will give you O(N^3).

To simplify the sum, I wrote the following Haskell program:

import Control.Monad
import Data.List
import qualified Data.Map as Map

-- Take equivalence relation, like [(1,2),(2,3)] and return its partition, like [3,1,1]
f xs = sort $ map length $ foldr f (map return [1..5]) xs
       where f (x,y) a = let [v1] = filter (x `elem`) a
                             [v2] = filter (y `elem`) a
                         in if v1 == v2 then a else (a \\ [v1,v2]) ++ [v1++v2]

-- All 1024 subsets of [(1,2),(1,3), ..., (4,5)]
subsets = filterM (const [False, True]) [(i,j) | i <- [1..5], j <- [i+1..5]]

res = Map.fromListWith (+) $ map (\k -> (f k, (-1)^(length k))) subsets

*Main> res
Loading package array-0.3.0.1 ... linking ... done.
Loading package containers-0.3.0.0 ... linking ... done.
fromList [([1,1,1,1,1],1),([1,1,1,2],-10),([1,1,3],20),([1,2,2],15),([1,4],-30),([2,3],-20),([5],24)]

which means that the formula is

all subsets - 10X₁ + 20X₂ - 30X₃ + 24X₄ + 15X₅ - 20X₆.

Check:

How many are there quintuples in [0,0,0,...,0] summing up to 0? One way to compute that is directly, second way is to use the formula (and not care about distinct positions):

direct x = x*(x-1)*(x-2)*(x-3)*(x-4)
indirect x = x^5 - 10 * x^4 + 20 * x^3 + 15 * x^3 - 30 * x^2 - 20*x^2 + 24*x

*Main> direct 100
9034502400
*Main> indirect 100
9034502400

Other remarks:

Also, there's O(a_n log a_n) solution: Compute (x^a₁ + ... + x^a_n)⁵ using FFT, the result is coefficient at x^S. This allows some a_i to be used twice, but you can subtract polynomials like (x^2a₁ + ... + x^2a_n)⁵*(x^a₁ + ... + x^a_n)³ etc. according to inclusion-exclusion principle.

In some restricted models of computation, it has been shown decision version of this problem requires O(N^3) time.