Give a natural positive Int (n) and a desired sum (k) find a random array with length n and sum k

Question

I want a function that given an N: Int and a K: Int returns an Array length N, elements from 1...N and total sum K. So for example: For n == 3 and k == 7 Solutions: [3, 1, 3], [3,3,1], [3,2,2] ... My naive solution is:

func arrayWith(n: Int, k : Int) -> [Int] {
    let elements = (1...n).map { $0 } // We only allow 1,2,3.... till n
    var output = [Int]()
    repeat {
        output.removeAll()
        for _ in 0..<n {
            let validNumbers = elements.filter { $0 + output.reduce(0,+) <= k }
            if let random = validNumbers.randomElement() {
                output.append(random)
            } else {
                break
            }
        }

    } while output.count != n || output.reduce(0,+) != k
    return output
}

This solution just works, but the time complexity is too high. arrayWith(n: 7, k: 49) (only a possibility [7,7,7,7,7,7,7], takes seconds.)

Note that N, so its not necessary to handle invalid cases.

Possibly helpful: [Split a random value into four that sum up to it](https://stackoverflow.com/questions/30154718/split-a-random-value-into-four-that-sum-up-to-it) – the same problem but without the restriction that each number is at most N. — Martin R, Apr 06 '20 at 09:21
What do you mean, "a random array"? Should all possible arrays be equally likely? What are you using this for? — Paul Hankin, Apr 06 '20 at 10:09
@PaulHankin the more random the better. What is used for? I'm just practicing old problems from Google Code Jam — Godfather, Apr 06 '20 at 10:52

score 1 · Accepted Answer · answered Apr 06 '20 at 10:11

Here is a very simple and fast solution producing a random array with the given sum and restrictions:

func arrayWith(n: Int, k: Int) -> [Int] {
    var output: [Int] = []
    var sum = k // Sum of the remaining elements
    for i in 1...n {
        let lo = max(sum - (n - i) * n, 1) // Minimum possible value for the next number
        let hi = min(sum - (n - i), n) // Maximum possible value for the next number
        let x = Int.random(in: lo...hi)
        sum -= x
        output.append(x)
    }
    return output
}

The idea is to choose each element randomly in the range 1...n, but also restricted in such a way that the given sum is still achievable with the remaining elements.

Note that this is a quite simple approach, and does not produce all possible solutions with equal probability. Shuffling the output array may improve that a bit.

For more sophisticated solutions have a look at Random numbers that add to 100: Matlab and Non biased return a list of n random positive numbers (>=0) so that their sum == total_sum where the problem is handled in other languages and in a theoretical, language-agnostic way.

Give a natural positive Int (n) and a desired sum (k) find a random array with length n and sum k

1 Answers1