[HashMap] HackerRank Algorithms in Java ‘Count Triplets’ solution

Problem

You are given an array and you need to find number of tripets of indices (i,j,k) such that the elements at those indices are in geometric progression for a given common ratio r and i < j < k. For example, arr = [1,4,16,64]. If r = 4, we have [1,4,16] and [4,16,64] at indices (0,1,2) and (1,2,3).

Solution

처음에 문제를 잘못 이해해서 고생했고, 나중에 이해를 한 뒤에는 시간초과로 고생을 했다. 해커랭크의 디스커션 discussions을 열심히 참고하고 열심히 이해한 끝에 다음과 같은 코드를 이해했다. (thanks to @beat1Percent)

• ‘geometric progression’에 주목해야 한다. 뜻은 등비수열이다. (A, B, C가 있을 때 B = A * r, C = B * r = A * r * r 이다. 따라서 A = B / r 이고 B = C / r 이다.)
• 예를 들어 arr = 1 3 9 9 27 81이고 r = 3일 때의 경우의 수는 [1,3,9], [1,3,9], [3,9,27], [3,9,27], [9,27,81], [9,27,81]로 6이 정답이다. 9를 보면 9가 끝에 있는 경우가 있고, 중간에 있는 경우가 있고, 처음에 있는 경우가 있다. 이 모든 경우를 생각하여 문제를 해결해야 한다.
• 현재 값이 처음에 왔을 경우의 수를 구하는 Map<Long, Long> mapForStart와 중간에 왔을 경우의 수를 구하는 Map<Long, Long> mapForMiddle를 둔다.
• 모든 숫자는 첫번째 위치에 올 수 있으므로 아무런 조건 없이 mapForStart.put(num, mapForStart.getOrDefault(num, 0L) + 1);를 수행한다.
• 현재 값을 r로 나눈 몫이(= $prev) 중간에 있었던 적이 있으면 (mapForMiddle.get(prev) != null) 현재 값까지 총 3가지의 숫자가 모인 것이므로 부분 등비수열을 완성했다. 지금까지의 경우의 수에 현재 값이 온 경우를 생각하는 것이므로 ans += mapForMiddle.get(prev)를 처리한다.
• 현재 값을 r로 나눈 몫이(= $prev) 처음에 있었던 적이 있으면 (mapForStart.get(prev) != null) 현재 값이 중간에 오는 경우의 수가 생기므로 mapForMiddle 맵에 그 경우의 수를 추가한다.
• 따라서 ans 값은 1 (1,3,9인 경우, 단 여기서 9는 첫번째 9) -> 2 (1,3,9인 경우, 단 여기서 9는 두번째 9) -> 4 (3,9,27인 경우, 이때 27이 마지막으로 오는 경우가 2번이기 때문) -> 6 (81이 마지막으로 오는 경우가 2번이기 때문) 으로 변한다. 답은 6이 된다.

Code

    // Complete the countTriplets function below.
    static long countTriplets(List<Long> arr, long r) {
        long ans = 0;
        
        Map<Long, Long> mapForStart = new HashMap<>();
        Map<Long, Long> mapForMiddle = new HashMap<>();
        for(long num : arr) {
            if(num % r == 0) {
                long prev = num / r;
                Long cntForEnd = mapForMiddle.get(prev); // num is the last number
                if(cntForEnd != null) ans += cntForEnd;
                
                Long cntForMiddle = mapForStart.get(prev); // num is the middle number
                if(cntForMiddle != null) mapForMiddle.put(num, mapForMiddle.getOrDefault(num, 0L) + cntForMiddle);
            }
            
            mapForStart.put(num, mapForStart.getOrDefault(num, 0L) + 1); // num is the first number
        }
        return ans;
    }