CODEWITHSUNDEEP
pythonlist-comprehension
S L
S L

Reputation: 14318

generating 3 numbers whose sum is n

I need to generate three natural numbers whose sum is n. The first number can be at max x, the second number can be at max y and the final number can be at max z. Currently I am doing this

def f(n):
    return [(i, j, k)
            for i in range(x+1)
            for j in range(y+1)
            for k in range(z+1)
            if i + j + k == n]

But n is very large, around 500 and x, y, z are less than 200. Currently I have 3 variables (i, j, k) generated from 3 ranges. Can this be done using two loops inside list comprehension?

Upvotes: 2

Views: 115

Answers (2)

Holt
Holt

Reputation: 37616

What you need to do is constrain the second variable (j) depending on the value of the first (i), and the last one (k) depending on the 2 firsts:

l = []
for i in range(1, x + 1):
    for j in range(max(1, n - x - z), min(y + 1, n - x - 1)): # No need to go above n - x - 1or bellow n - x - z
        # For k, you do not even need a loop since it is constrained by i and j
        k = n - i - j 
        if 0 < k <= z:
            l.append((i, j, k))

I am not using list compression to details the code, but of course you can:

[(i, j, n - i - j) for i in range(1, x + 1) for j in range(max(1, n - x - z), min(y + 1, n - x - 1)) if 0 < n - i - j <= z]

When x equals 1 (first value in the global loop), the range for y goes from n - x - z = 500 - 1 - 200 = 299 to 200, i.e. you do not even enter the y loop, and that's normal since you need x to be at least 100 to be able to reach 500 (because y and z are at most 200).

Every time you are facing this kind of problem, think about « How can you reduce the domain of the current variable depending on the already assigned variables? ».

By the way, doing the k = n - i - j assignment is equivalent of looping on a range of size 0 or 1 like so:

for k in range(max(1, n - i - j), min(z + 1, n - i - j + 1)):
    l.append((i, j, k))

Notice that in this case, you do not have to check for 0 < k <= z anymore.

Upvotes: 2

Sorin
Sorin

Reputation: 11968

Yup,

You can just compute the third number and check that it's in the right range (Martijn Pieters♦).

[(i, j, n - i - j) for i in range(x+1) for j in range(y+1) if 0 <= n - i - j <= z]

Upvotes: 3

Related Questions