Find all triplets i,j,k such that i+j+k=n

Viewed 853

I've coded this but this is very long:

for i in range(n + 1):
    for j in range(n + 1):
        for k in range(n + 1):
            if i + j + k == n:

Is there a clever way to make it go faster? Currently it's O(n^3) which is quite sad.

3 Answers

There are a few possible solutions - you don't actually need to loop till N in all of them, and the last number comes completely free.

Keep in mind all numbers in triplet must be positive (otherwise the answer is infinite).

  1. If permutations are not included (i.e. (1,2,3) is the same triplet as (3,2,1)), going from smallest to largest:

    def iter_triplets(n):
        # This is the smallest number, can't be more than 1/3 of n
        for i in range(0, n//3 + 1):
            sum_left = n-i
            # This is the second smallest, can't be more than 1/2 of the sum_left or less than the first by definition
            for j in range(i, sum_left//2 + 1):
                yield (i, j, sum_left-j)  # Last number is calculated.
    
    
    >>> list(iter_triplets(6))  
    [(0, 0, 6), (0, 1, 5), (0, 2, 4), (0, 3, 3), (1, 1, 4), (1, 2, 3), (2, 2, 2)]
    >>> list(iter_triplets(10))
    [(0, 0, 10), (0, 1, 9), (0, 2, 8), (0, 3, 7), (0, 4, 6), (0, 5, 5), (1, 1, 8), (1, 2, 7), (1, 3, 6), (1, 4, 5), (2, 2, 6), (2, 3, 5), (2, 4, 4), (3, 3, 4)]
    
  2. If permutations are not included (i.e. (1,2,3) is the same triplet as (3,2,1)), going from largest to smallest:

    import math
    def iter_triplets(n):
        # This is the biggest number, can't be less than 1/3 of n
        for i in range(n, math.ceil(n/3) - 1, -1):
            sum_left = n-i
            # This is the second biggest number, can't be less than 1/2 of the sum_left and more than first number by definition.
            # ceil to correct rounding errors.
            for j in range(min(sum_left, i), math.ceil((sum_left)/2) - 1, -1):
                yield (i, j, sum_left-j)  # Last number is calculated.
    
    
    >>> list(iter_triplets(10))
    [(10, 0, 0), (9, 1, 0), (8, 2, 0), (8, 1, 1), (7, 3, 0), (7, 2, 1), (6, 4, 0), (6, 3, 1), (6, 2, 2), (5, 5, 0), (5, 4, 1), (5, 3, 2), (4, 4, 2), (4, 3, 3)]
    
  3. If permutations are included (i.e. (1,2,3) is a different triplet than (3,2,1)), going from smallest to largest:

    def iter_triplet_permutations(n):
        for i in range(0, n+1):
            sum_left = n-i
            for j in range(0, sum_left+1):
                yield (i, j, sum_left-j)
    
    >>> list(iter_triplet_permutations(5)) 
    [(0, 0, 5), (0, 1, 4), (0, 2, 3), (0, 3, 2), (0, 4, 1), (0, 5, 0), (1, 0, 4), (1, 1, 3), (1, 2, 2), (1, 3, 1), (1, 4, 0), (2, 0, 3), (2, 1, 2), (2, 2, 1), (2, 3, 0), (3, 0, 2), (3, 1, 1), (3, 2, 0), (4, 0, 1), (4, 1, 0), (5, 0, 0)]
    

The innermost loop seems redundant since once you have i and j, k comes for free:

for i in range(n + 1):
    for j in range(n + 1):
        if i+j <= n:
            print((i, j, n-i-j))

If we want triplets of unique numbers that add to n, then maybe this could work:

for i in range(n + 1):
    for j in range(i, (n - i)//2 + 1):
        out.append((i, j, n-i-j))

There are few ways to make the code run faster.

  1. Use a list comprehension. Syntax [ReturnThing ForLoops Condition]

  2. Start each loop from where the previous loop got to. This avoid duplicates

    def MakeTriplet(n):                                                                          
        return [(i,j,k) for i in range(0,n+1) for j in range(i,n+1) for k in range(j,n+1) if (i+j+k)==n]
    
Related