HackerRank Count Triplets

아이디어를 얻기 어려워 검색을 해봤다
아이디어 출처

지정된 값 v를 기준으로 왼쪽에는 v를 r로 나눈 몫이 key로 들어있는 dict, 오른쪽에는 v에 r을 곱한 값이 key로 들어있는 dict를 만들어야 한다.

for문 돌 때마다 만들면 timeout이 나므로 초기에 기준점 이후의 after_dict를 미리 만들어놓고 v를 옮겨가면서 before_dict, after_dict에서 적절한 값을 찾아 개수를 곱해 누적해주면 가능한 경우의 수를 구할 수 있다.

#!/bin/python3

import os
from collections import Counter


def count_triplets(arr, r):
    count = 0
    after_dict = Counter(arr)
    before_dict = {}

    for v in arr:
        before = v // r
        after = v * r
        after_dict[v] -= 1  # exclude current value
        if v % r == 0 and before in before_dict and after in after_dict:
            count += before_dict[before] * after_dict[after]
        before_dict[v] = before_dict[v] + 1 if v in before_dict else 1

    return count


if __name__ == '__main__':
    fptr = open(os.environ['OUTPUT_PATH'], 'w')

    nr = input().rstrip().split()

    n = int(nr[0])

    r = int(nr[1])

    arr = list(map(int, input().rstrip().split()))

    ans = count_triplets(arr, r)

    fptr.write(str(ans) + '\n')

    fptr.close()