Find all Triplets with zero sum

Problem Statement

Given an array of unsorted numbers, find all unique triplets in the array whose sum is zero. The array may have duplicates.

Example 1:

Input: a[] = [-5, 3, 2, -3, 1]
Output: [-5, 3, 2] [2, -3, 1]

Example 2:

Input: a[] = [-1, -1, 2, -1, -1, 2, -1, -1, 2]
Output: [-1, -1, 2]

Solution

Two pointer approach

import java.util.ArrayList;
import java.util.Arrays;
import java.util.List;
import java.util.Scanner;

public class TripletsWithZeroSum {
    private static List<List<Integer>> findTripletsWithZeroSum(int[] a) {
        List<List<Integer>> triplets = new ArrayList<>();

        Arrays.sort(a);
        for(int i = 0; i < a.length-2; i++) {
            // Skip duplicates
            if(i > 0 && a[i] == a[i-1]) {
                continue;
            }

            // Fix one number
            int firstNum = a[i];

            // Use Two-sum approach to get the other two numbers
            int left = i+1, right = a.length-1;
            while(left < right) {
                int sum = firstNum + a[left] + a[right];
                if(sum == 0) {
                    triplets.add(Arrays.asList(firstNum, a[left], a[right]));
                    left++;
                    right--;

                    // Skip duplicates
                    while(left < right && a[left] == a[left-1]) {
                        left++;
                    }
                    while(left < right && a[right] == a[right+1]) {
                        right--;
                    }
                } else if(sum < 0) {
                    left++;
                } else {
                    right--;
                }
            }
        }

        return triplets;
    }

    public static void main(String[] args) {
        Scanner keyboard = new Scanner(System.in);
        int n = keyboard.nextInt();
        int[] a = new int[n];
        for(int i = 0; i < n; i++) {
            a[i] = keyboard.nextInt();
        }
        keyboard.close();

        List<List<Integer>> triplets = findTripletsWithZeroSum(a);
        for(List<Integer> triplet: triplets) {
            for(Integer num: triplet) {
                System.out.print(num + " ");
            }
            System.out.println();
        }
    }
}

Simplifying the code by extracting out the two-sum code into a separate function

import java.util.ArrayList;
import java.util.Arrays;
import java.util.List;
import java.util.Scanner;

public class TripletsWithZeroSum {
    private static List<List<Integer>> findTripletsWithZeroSum(int[] a) {
        List<List<Integer>> triplets = new ArrayList<>();

        Arrays.sort(a);
        for(int i = 0; i < a.length-2; i++) {
            // Skip duplicates
            if(i > 0 && a[i] == a[i-1]) {
                continue;
            }
            // Fix one number a[i] and find pairs with sum -a[i] starting from index i+1
            addPairsWithTargetSum(a, -a[i], i+1, triplets);
        }

        return triplets;
    }

    private static void addPairsWithTargetSum(int[] a, int targetSum, int left, List<List<Integer>> triplets) {
        int right = a.length-1;

        while(left < right) {
            int sum = a[left] + a[right];
            if(sum == targetSum) {
                triplets.add(Arrays.asList(-targetSum, a[left], a[right]));
                left++;
                right--;

                // Skip duplicates
                while(left < right && a[left] == a[left-1]) {
                    left++;
                }
                while(left < right && a[right] == a[right+1]) {
                    right--;
                }
            } else if(sum < targetSum) {
                left++;
            } else {
                right--;
            }
        }
    }
}