Count Triplets

def countTriplets(arr, r):
    m2 = defaultdict(int)
    m3 = defaultdict(int)

    total = 0
    for val in arr:
        if val in m3:
            total += m3[val]
        # 1 2 2 4
        if val in m2:
            m3[val * r] += m2[val]

        m2[val * r] += 1

    return total

 

이 문제를 처음 풀었을 때는 삼중 포문을 사용 했었다. 첫번째 값을 기준으로 두번째, 세번째 값을 찾는 방식이었는데 세 숫자의 인덱스가 a < b < c 순이 되는지 확인했다.

 

그 상태에서 두번째, 세번째 숫자를 찾는 방법을 아무리 최적화 하려고 해도 방법이 안떠오르는 것이다.

그래서 첫번째 값을 기준으로 하는게 아니라 가운데 값을 기준으로 하는 방식으로 수정했다.

첫번째 값은 중요치 않고, 두번째 & 세번째 값을 저장하는게 핵심이다.

로직은 간단하다.

1. 두 개의 딕셔너리를 만든다.

   1. 하나는 가운데 값 을 저장
   2. 하나는 세번째 값 을 저장

2. 세번째 값 딕셔너리에 현재 값이 있으면 카운팅 +1

3. 두번째 값 딕셔너리에 현재 값이 있으면 두번째 값 딕셔너리에 있는 값을 세번째 값 딕셔너리 에 넣어주는데 그 때 현재 값 * r 을 키 값으로 넣어준다.

4. 두번째 값 딕셔너리 에 현재 값 *r 을 키로 하고 카운팅 +1 을 해준다.