Triple sum

def triplets(a, b, c):
    a = list(set(a))
    b = list(set(b))
    c = list(set(c))

    count = 0
    a.sort()
    b.sort()
    c.sort()

    frontIdx = 0
    backIdx = 0

    for ct in b:
        while frontIdx < len(a) and a[frontIdx] <= ct:
            frontIdx += 1

        while backIdx < len(c) and c[backIdx] <= ct:
            backIdx += 1
        
        count += (frontIdx * backIdx)

    return count

 

먼저 a,b,c에 중복 항목이 없는 리스트로 변환하고 정렬을 한다.

이후에는 기준이 되는 b를 순회하면서 앞,뒤 값을 비교하면 된다. 여기서 포인트는 정렬이다.

정렬을 하게 되면 보장되는 것이 있다.

 

기준이 되는 값(b 리스트)보다 작은 값이 앞,뒤 리스트에 존재하고 끝 인덱스까지 도달하면 b 리스트의 나머지 값들은 자동으로 크다는 것이 보장된다. 그래서 앞,뒤 리스트를 순회하면서 해당 인덱스를 저장해놓으면 순회하는 횟수가 많이 줄어든다.

 

이진 탐색을 사용한 해결책

def triplets(a, b, c):
    a, b, c = sorted(set(a)), sorted(set(b)), sorted(set(c))
    return sum((bisect(a, x) * bisect(c, x) for x in b))

 

문제를 봤을 때 이진탐색을 이용해서 하면 되겠군 생각을 했었는데 bisect 모듈의 사용법이 익숙치 않아서 제대로 결과값이 안나왔었다.

위의 예제를 보고 도움이 많이 됐다.