amazon music interview question solution in python

amazon_music_pair.py

import itertools


def music_pair(song_durations, travel_duration):
    same_duration_song_indexes = {}
    for i, duration in enumerate(song_durations):
        try:
            same_duration_song_indexes[duration].append(i)
        except KeyError:
            same_duration_song_indexes[duration] = [i]
    song_durations.sort()
    l = 0
    h = len(song_durations) - 1
    pair_durations = []
    while l < h:
        pair_duration = song_durations[l] + song_durations[h]
        if pair_duration == travel_duration - 30:
            pair_durations.append((song_durations[l], song_durations[h]))
            l += 1
            h -= 1
        elif pair_duration > travel_duration - 30:
            h -= 1
        else:
            l += 1
    output = []
    for pair_duration in pair_durations:
        output.extend(list(itertools.product(same_duration_song_indexes[pair_duration[0]],same_duration_song_indexes[pair_duration[1]])))
    return output