patternjavaModerate

Finding Pythagorean triplet in array

Submitted by: @import:stackexchange-codereview·Mar 10, 2026·

Viewed 0 times

pythagoreantripletarrayfinding

Problem

We have an integer array as:

private int[] arr = {1, 3, 5, 14, 18, 29, 78};

We have a function which takes three inputs of the array and checks whether:

a * a = b * b + c * c

If the function returns true then those 3 inputs are stored in an ArrayList:

private boolean findTriplet(int a, int b, int c) {
    int squareA = a * a;
    int squareB = b * b;
    int squareC = c * c;
    int sum = squareB + squareC;
    if (squareA == sum) {
        return true;
    }
    return false;
}

Iterating through the array:

private void getTriplets(int[] arr) {
    for (int i = 0; i < arr.length; i++) {
        for (int j = 1; j < arr.length; j++) {
            if (i != j) {
                for (int k = i+j; k < arr.length; k++) {
                    if ((i != k) || (j != k)) {
                        boolean tripResult = findTriplet(arr[i], arr[j], arr[k]);
                        if (tripResult) {
                            tripList.add(new Triplet(arr[i], arr[j], arr[k]));
                        }
                    }
                }
            }
        }
    }
}

According to the above solution the complexity of the program is \$O(n^3)\$.

How can I optimise this solution or get a more flexible complexity?

Solution

A couple of notes on your code before going into alternative solutions.

-
In your findTriplet method, you are doing:

if (squareA == sum) {
    return true;
}
return false;

which is more concisely expressed as

return squareA == sum;

-
Your main getTriplets method checks whether (i != k) || (j != k), but this isn't possible by construction: k starts at i+j so it cannot be equal to i or j; thus this check can be removed. Also j starts at 1 when you could make it start at j + 1, which would also get rid of the i != j check.

You can do this in O(n²) by utilizing a TreeSet.

The idea is the following:

Transform the array into a TreeSet, where each element will be the squared value of the element in the array. Using a TreeSet has the advantage that it sorts the elements ascendingly and allows for fast look-up.

For each element in the set, get all the elements greater than it and see if the set contains a third value equal to the sum of those two. This is because, if we look for 3 numbers a, b and c where a + b = c, then having both a and b only leaves the question: does the set contains a + b?

A possible implementation is the following:

private void getTriplets(int[] arr) {
    NavigableSet set = new TreeSet<>(); 
    for (int element : arr) {
        set.add(element * element);
    }
    for (Integer a : set) {
        for (Integer b : set.tailSet(a, false)) {
            if (set.tailSet(b, false).contains(a + b)) {
                tripList.add(new Triplet(a, b, a + b));
            }
        }
    }
}

The method tailSet(element, false) will return a view of the set with all the elements strictly greater than the passed element. This method runs in constant-time. In the code above, a will loop through all the values of the set, b will loop over the values greater than a and we see if the set contains a + b. If it does then we have found a Pythagorian triplet.

This solution does have the overhead of creating a Set (so O(n) memory) but it is simple to write.

Another approach with O(1) memory is to keep an array but the approach is very similar:

Sort the array in ascending order. This takes O(n log n).

-
Now consider each element a[i]. If a[i]=a[j]+a[k], then, since numbers are positive and array is now sorted,

k
  
Repeat step 2 for each index i. This way you'll need totally O(n²) operations, which will be the final estimate.


  

Step 2 is similar to finding a given number by summing up integers from 2 sorted arrays.

In terms of code, you could have the following:

private void getTriplets(int[] arr) {
    int[] squaredArray = new int[arr.length];
    for (int i = 0; i = 0; i--) {
        int j = 1, k = i;
        while (j  squaredArray[i]) {
                k--;
            } else {
                tripList.add(new Triplet(squaredArray[j], squaredArray[k], squaredArray[i]));
                j++;
                k--;
            }
        }
    }
}


The main loop on

i is performed descendingly since we're looking for two elements whose sum is squaredArray[i]` (and the array is sorted ascendingly).

Code Snippets

if (squareA == sum) {
    return true;
}
return false;

return squareA == sum;

private void getTriplets(int[] arr) {
    NavigableSet<Integer> set = new TreeSet<>(); 
    for (int element : arr) {
        set.add(element * element);
    }
    for (Integer a : set) {
        for (Integer b : set.tailSet(a, false)) {
            if (set.tailSet(b, false).contains(a + b)) {
                tripList.add(new Triplet(a, b, a + b));
            }
        }
    }
}

private void getTriplets(int[] arr) {
    int[] squaredArray = new int[arr.length];
    for (int i = 0; i < arr.length; i++) {
        squaredArray[i] = arr[i] * arr[i];
    }
    Arrays.sort(squaredArray);
    for (int i = arr.length - 1; i >= 0; i--) {
        int j = 1, k = i;
        while (j < k) {
            int sum = squaredArray[j] + squaredArray[k];
            if (sum < squaredArray[i]) {
                j++;
            } else if (sum > squaredArray[i]) {
                k--;
            } else {
                tripList.add(new Triplet(squaredArray[j], squaredArray[k], squaredArray[i]));
                j++;
                k--;
            }
        }
    }
}

Context

StackExchange Code Review Q#131701, answer score: 12

Revisions (0)

No revisions yet.