[Hackrrank] [Medium] Count Triplets

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

Count Triplets | HackerRank

풀이

• 숫자로 구성된 배열, 등비급수 공비 r이 주어졌을 때, 세 개의 숫자의 index가 오름차순을 만족하는 배열의 수를 구하는 문제
• 단순히 숫자 종류만 count해서 곱하면 정답보다 많은 수를 계산하게되므로, index가 오름차순인 조건을 만족하는게 중요
• 단순반복으로 구현하면 timeout 발생.
• dictionary 자료구조를 사용해 현재 숫자에서 가능한 등비수열 개수를 count하는 dynamic programming 문제
• 세 가지 경우를 생각해야 함. 등비수열이 오름차순으로 num1, num2, num3이라 가정. arr 특정 원소는 num이라 가정.
  • num이 num1일 때(가장 작은수) -> num2에 num*r(가운데 수)을 key값으로 한 뒤 1 증가
  • num이 num2일 때 -> num3에 key값으로 num * r(가장 큰 수)을 하고 num2[num] 증가
    • num3[num * r] += num2[num]
  • num이 num3일 때 -> num3 dictionary에는 현재까지 본 데이터로 num이 가장 큰 수일 때 등비수열 개수가 저장되어 있고 이를 result에 더해줌.
  • num이 num3에도, num2에도 포함되는 경우도 있음. -> if문으로 처리(not elif)
    • r=3일 때, 9같은 경우 1, 3, 9(가장 큰 수) // 3, 9, 27(중간 숫자)
• 순서대로 iteration 돌기 때문에 index의 오름차순 조건을 만족.

def countTriplets(arr, r):
    from collections import defaultdict
    result = 0
    num2 = defaultdict(int)
    num3 = defaultdict(int)

    for num in arr:
        if num in num3:
            result += num3[num]
        if num in num2:
            num3[num * r] += num2[num]
        num2[num * r] += 1 
    return result