Pairs of elements

You are given an array of length \(N\). You are required to count the number of \((i, j)\) pairs where \(1⩽i<j⩽N\) such that the difference of the array elements on that indices is equal to the sum of the square of their indices.

That is the count of the number of pairs of \((i, j)\) such that it satisfies this equation (\(A[j] - A[i] = i^2 + j^2\)). 

Example:

Input:  n=5, a[]={4, 9, 6, 29, 30}
Output: 3

Approach

Java

import java.util.ArrayList;
import java.util.HashMap;

public class PairsOfElements {

    public static void main(String[] args) {
        int n = 5;
        long[] a = { 4, 9, 6, 29, 30 };
        HashMap<Long, ArrayList<Integer>> hash = new HashMap<>();
        HashMap<Long, ArrayList<Integer>> map = new HashMap<>();

        for (int i = 0; i < n; i++) {
            long v1 = ((long) (i + 1) * (long) (i + 1)) + a[i];
            long v2 = a[i] - ((long) (i + 1) * (long) (i + 1));
            if (hash.containsKey(v1)) {
                ArrayList<Integer> al = hash.get(v1);
                al.add(i);
                hash.put(v1, al);
            } else {
                ArrayList<Integer> al = new ArrayList<>();
                al.add(i);
                hash.put(v1, al);
            }
        }
        long count = 0;
        for (int j = 1; j < n; j++) {
            long v1 = a[j] - ((long) (j + 1) * (long) (j + 1));
            // long v2=a[j]+(j+1)*(j+1);
            if (hash.containsKey(v1)) {
                ArrayList<Integer> al = hash.get(v1);
                count += al.size();
            }
        }

        System.out.println(count);

    }

    static class Pair {
        public final int first;
        public final int second;

        public Pair(int a, int b) {
            first = a;
            second = b;
        }

        public boolean equals(Object o) {
            if (!(o instanceof Pair))
                return false;
            Pair p = (Pair) o;
            return first == p.first && second == p.second;
        }

        public int hashCode() {
            return first + second << 16;
        }

    }
}