[Hackerrank] Triple sum

https://www.hackerrank.com/challenges/triple-sum/problem

Triple sum | HackerRank

풀이

• input : 숫자로 구성된 리스트 3개 (a, b, c)
• a, c의 원소는 b의 원소보다 작은 조건으로 생성할 수 있는 tuple의 개수를 구하는 문제
• b의 i번째 원소 bi일 때 len([bi보다 작은 a의 원소]) * len([bi보다 작은 c의 원소])를 구해서 모든 원소 b에 대해 1번 순회하면 됨.
  • [bi보다 작은 a의 원소], [bi보다 작은 c의 원소]를 list comprehension으로 구현하면 당연히 timeout 발생.
• 문제 조건에 a, b, c 리스트는 정렬된 순서라고 명시되어있지 않고, 중복으로 입력될 수 있음.
  • 보통 중복을 제거할 때 set( )을 사용해 list로 감싸줌. 그리고sorted( )를 사용하면 빠르게 정렬 가능.
• 입력으로 주는 3개의 리스트 전부 정렬해줬기 때문에, 각각의 index를 순회하면서 b의 원소보다 클 때의 index를 구해주면 됨.
• 시간복잡도 O(NlogN)으로 구현.

def triplets(a, b, c):
    # a <= b >= c
    a = sorted(list(set(a)))
    b = sorted(list(set(b)))
    c = sorted(list(set(c)))
    counts = 0
    index_a = 0
    index_c = 0
    for index_b in range(len(b)):
        while index_a < len(a) and a[index_a] <= b[index_b]:
            index_a += 1

        while index_c < len(c) and c[index_c] <= b[index_b]:
            index_c += 1

        counts += index_a * index_c
    return counts