beautiful_triplets.py

# source: https://www.hackerrank.com/challenges/beautiful-triplets/problem
# video: https://youtu.be/fOUWTyxYXUQ

def beautifulTriplets(d, arr):
    # Keep a running count dict
    counts = {}
    total = 0
    # loop through list to get counts
    for val in arr: # O(n)
        if val not in counts:
            counts[val] = 0
        counts[val] += 1
        # check for triplet
        first_val_count = counts.get(val - d)
        second_val_count = counts.get(val - d - d)
        if first_val_count and second_val_count:
            total += first_val_count * second_val_count
    
    return total # Total Time Complexity = O(n)