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

Question

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.

Answer 1

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)]