Count Triplets

  • + 1 comment

    Eritorial transformed into Java. Don't forget to check if number is divisible by r.

        // Complete the countTriplets function below.
        static long countTriplets(List<Long> arr, long r) {
            HashMap<Long, Long> right = new HashMap<>();
            HashMap<Long, Long> left = new HashMap<>();
            
            for (Long num : arr) {
                right.put(num, right.getOrDefault(num, 0L) + 1);
            }
            
            long tripletsCounter = 0;
            
            for (Long num : arr) {
                if (right.getOrDefault(num, 0L) > 0) {
                    right.put(num, right.get(num) - 1);
                }
                if (num % r == 0) {
                    tripletsCounter += right.getOrDefault(num * r, 0L) * left.getOrDefault(num / r, 0L);
                }
                left.put(num, left.getOrDefault(num, 0L) + 1);
            }
            
            return tripletsCounter;
        }