UNO.ai: Triplets

unoai_main_2.py

def get_triplets(arr, to_sum):
    
    # Array length.
    n = len(arr)
    
    # Iterate through the array twice [O(n^2)] and check if the remainder of the sum is present in the dict.
    # This helps bring the algorithm from O(n^3) to O(n^2).
    ans = []
    for i in range(n - 2):
        pseudo_sums = dict()
        for j in range(i + 1, n):
            x = to_sum - (arr[i] + arr[j])
            
            # Check is remainder of sum is present, if yes then add triplet to ans.
            if x in pseudo_sums.keys():
                ans.append([x, arr[i], arr[j]])
            else:
                pseudo_sums[arr[j]] = 1
    return ans

# Initialize elements.
arr = [0, 2, 3, 4, 5, 7, 8]
to_sum = 7

# Call main function.
ans = get_triplets(arr, to_sum)
print("\nThe triplets that sum to 7 are:\n", ans)