HackerRank : Count Triplets

문제 링크

The problem link : Count Triplets

문제 요약

배열의 사이즈 n(1<=n<=10^5), 공비의 크기 r(1<=r<=10^9), 그리고 n개의 숫자들(1<=arr[i]<=10^9)이 나열됩니다.

4 2 // n, r
1 2 2 4 // n개의 숫자

arr배열에서 주어진 공비(r)를 만족하는 등비수열의 개수를 구하세요.

예시로, 위의 배열에서 주어진 조건을 만족하는 등비수열의 개수는 다음과 같습니다.

정답 : 2
arr[0]=1, arr[1]=2, arr[3]=4 -> {0,1,3}
arr[0]=1, arr[2]=2, arr[3]=4 -> {0,2,3}

문제 풀이

핵심 코드는 다음과 같습니다.

    static long countTriplets(List<Long> arr, long r) {
        long ans = 0L;
        HashMap<Long, Long> right = new HashMap<>(); // 1
        HashMap<Long, Long> left = new HashMap<>();

        for(Long num : arr) { // 2
            right.put(num, right.getOrDefault(num, 0L)+1);
        }

        for(Long num : arr) {
            if(right.get(num) > 0){ // 3
                right.put(num, right.get(num)-1);
            }
            if(num % r == 0) { // 4
                ans += right.getOrDefault(num*r, 0L) * left.getOrDefault(num/r, 0L);
            }
            left.put(num, left.getOrDefault(num, 0L)+1); // 5
        }
        return ans;
    }

(1) HashMap으로 right와 left를 선언합니다.
(2) right(=등비수열을 만족하는지 확인하기 위한 현재 숫자)에 arr에 있는 값들의 개수를 저장합니다.
(3) arr로 for문을 돌렸을때, 현재 위치의 숫자가 right에 저장되어있을 경우, 우선 개수를 -1 해줍니다.
(4) 만약 현재 숫자(=num)이 공비(=r)의 배수일 경우...
-> 정답에 right.getOrDefault(num*r, 0L) + left.getOrDefault(num/r, 0L)을 더해줍니다.
-> right.getOrDefault(num*r, 0L)의 뜻은, num의 바로 다음 공비의 배수가 있는지 판단하는 식입니다.
-> left.getOrDefault(num/r, 0L)의 뜻은, num의 바로 이전 공비의 배수가 있는지 판단하는 식입니다.
(5) 그렇게 탐색이 끝난 현재의 숫자(=num)은 left에 저장해줍니다.

시간 복잡도

O(n)입니다.