[HackerRank][Python3] Count Triplets

Problem :

https://www.hackerrank.com/challenges/count-triplets-1/problem

My Solution :

#!/usr/bin/env python3

from collections import defaultdict


def countTriplets(arr, r):
    total = 0
    single, double = defaultdict(int), defaultdict(int)
    for a in arr:
        if a % r == 0:
            total += double[a//r]
            double[a] += single[a//r]
        single[a] += 1
    return total


n, r = map(int, input().split())
arr = list(map(int, input().rstrip().split()))
ans = countTriplets(arr, r)
print(ans)

My Solution2 :

위 DP 방식은 원리를 바로 이해하기 어려울 수 있다. 아래 방식은 조금 더 직관적이다.

#!/usr/bin/env python3

from collections import defaultdict


def countTriplets(arr, r):
    total = 0
    small, large = defaultdict(int), defaultdict(int)
    for a in arr:
        large[a] += 1
    for a in arr:
        large[a] -= 1
        if a % r == 0:
            total += large[a*r] * small[a//r]
        small[a] += 1
    return total


n, r = map(int, input().split())
arr = list(map(int, input().rstrip().split()))
ans = countTriplets(arr, r)
print(ans)