Finding Pythagorean triplet in array

Question

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?

Answer 1

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(n2) by utilizing a TreeSet.

The idea is the following:

1. Transform the array into a TreeSet<Integer>, 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.
2. 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<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));
 }
 }
 }
}

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:

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

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

   To find such indexes, run a loop that increases j from 1 to i, and decreases k from i to 0 at the same time, until they meet. Increase j if a[j]+a[k] < a[i], and decrease k if the sum is greater than a[i]. If the sum is equal, that's one of the answers, print it, and shift both indexes.

   This takes O(i) operations.

3. Repeat step 2 for each index i. This way you'll need totally O(n2) 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 < 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--;
 }
 }
 }
}

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).

Answer 2

You tagged the question with "array", but I think this is the wrong data structure if you want speed: You need to cycle only over a and b, if you can find c2 quickly, and this is difficult with arrays, while Maps or Sets provide typically O(log n) access time.

I used a TreeMap here in order to being able to look up the squares quickly, but to recover the original values as well. Note that having the tailMap method is crucial for avoiding overhead and duplicates.

private static void getTriplets(int[] arr) {
 SortedMap<Integer, Integer> sqr = new TreeMap<>();
 Arrays.stream(arr).forEach(n -> sqr.put(n * n, n));
 sqr.forEach((aa, a) -> sqr.tailMap(aa).forEach(
 (bb, b) -> Optional.ofNullable(sqr.get(aa + bb)).ifPresent(
 c -> tripList.add(new Triplet(a, b, c)))
 ));
}

• Tunaki: Using forEach is generally not a good idea: it is preferable to use a reduce / collect approach that parallelize well. For example, in the first case, it should be Arrays.stream(arr).boxed().collect(Collectors.toMap(n -> n * n, n -> n));
• Landei: That won't work here, as I need a SortedMap. Too bad there is no Collector for this.
• Tunaki: Yes, there is: docs.oracle.com/javase/8/docs/api/java/util/stream/…. You can give a mapSupplier. Unfortunately, you also need to give a mergeFunction that does nothing since there won't be collisions here... See also stackoverflow.com/questions/31004899/…

Answer 3

Others have proposed some alternatives to your original algorithm and data structure. I'm going to discuss ways to optimize it instead.

findTriplet

private boolean findTriplet(int a, int b, int c) {

I'd prefer

private boolean isTriplet(int a, int b, int c) {

That's a more common name for something that tests for truth. And this method doesn't search for anything.

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

If the array is sorted ascending, then c will be greater than either a or b. So

 int sum = squareB + squareA;
 if (squareC == sum) {

would make more sense.

 int squareA = a * a;
 int squareB = b * b;
 int squareC = c * c;
 int sum = squareA + squareB;
 return sum == squareC;

The if is unnecessary.

 int squareA = a * a;
 int squareB = b * b;
 int squareC = c * c;
 return squareA + squareB == squareC;

This is simpler.

 return a * a + b * b == c * c;

Simplest.

getTriplets

private void getTriplets(int[] arr) {

I'd probably make this

private void findTriplets(int[] candidates) {

I find candidates to be more descriptive than arr. Another option would be sides.

 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)) {

The ((i != k) || (j != k)) should be ((i != k) && (j != k)). I think that k = i+j is also wrong. Perhaps you wanted k = 1 + Math.max(i, j) instead.

There are a couple ways to simplify this. In general, rather than nesting if statements, you can instead do

 for (int i = 0; i < candidates.length; i++) {
 for (int j = 1; j < candidates.length; j++) {
 if (i == j) {
 continue;
 }
 for (int k = 1 + Math.max(i, j); k < candidates.length; k++) {
 if ((i == k) || (j == k)) {
 continue;
 }

Note that the || is now the correct operator, as we took the logical negative of the original expression.

But in this case we have a simpler option.

 for (int i = 0; i < candidates.length; i++) {
 for (int j = i + 1; j < candidates.length; j++) {
 for (int k = j + 1; k < candidates.length; k++) {

This will check each combination of three elements of the array. If the array is sorted ascending, this will be sufficient. It never has the indexes the same, so you don't have to do separate equality checks.

 boolean tripResult = findTriplet(arr[i], arr[j], arr[k]);
 if (tripResult) {

You could simplify this to just

 if (isTriplet(candidates[i], candidates[j], candidates[k])) {

But we could do even better if the array is sorted.

 int comparison = compareTriplet(candidates[i], candidates[j], candidates[k]);
 if (comparison == 0) {
 tripList.add(new Triplet(candidates[i], candidates[j], candidates[k]));
 }
 if (comparison <= 0) {
 break;
 }

where

private boolean compareTriplet(int a, int b, int c) {
 return Integer.compare(a * a + b * b, c * c);
}

If the comparison is greater than 0, we can continue.

If the comparison is equal to 0, then we can add the triplet to the list.

If the comparison is less than or equal to 0, then we don't need to continue checking values for the third member of the triplet. We've either found or passed the point where we could have found the third member.