Counting triplets with given sum

Given an array of distinct numbers, how do we count the number of triplets with given sum?

For example let us consider the array {5, 2, 7, 8, 4, 3} and the sum is 14,
There are 3 triplets with the given sum
(2,5,7)
(2,4,8)
(3,4,7)

A brute-force algorithm will take O(n³) time. But this problem can also be solved based on 2 Sum problem. Here is the approach.

Sort the array first.
Fix the first element in the triplet, and search for a pair of elements whose sum is (Sum – Fixed element)
Repeat the above step for all the elements

We are essentially searching for a pair of elements for each fixed element. So the time complexity of this solution would be O(n²).

One variation/follow-up of this problem could be to find out the number of triplets with sum less than or equal to the given value.

Considering the first example again, there are 8 triplets with sum less than or equal to 14.

2 3 4
2 3 5
2 3 7
2 3 8
2 4 5
2 4 7
3 4 5
3 4 7

The solution for this is similar to the first problem and uses the trick discussed in my previous post.

The code for both is given in the following Java program.

	import java.util.Arrays;

	/**
	* Created by Ravi on 04-Dec-14.
	*/
	public class KSumTripletCount
	{
	public static void main(String[] args)
	{
	int [] array1 = {5, 2, 7, 8, 4, 3};
	System.out.println(getTripletWithSum(array1, 14));
	System.out.println(getTripletWithSumOrLess(array1, 14));
	}
	public static int getTripletWithSum(int[] array, int sum)
	{
	Arrays.sort(array);
	int tripletCount = 0;
	for(int i = 0; i < array.length-2; i++ )
	{
	for(int j = i+1, k = array.length-1; j < k; )
	{
	int curSum = array[i] + array[j] + array[k];
	if( curSum == sum )
	{
	System.out.println(array[i] + " " + array[j] + " " + array[k]);
	j++;
	k--;
	tripletCount++;
	}
	else if( curSum < sum )
	{
	j++;
	}
	else
	{
	k--;
	}
	}
	}
	return tripletCount;
	}

	public static int getTripletWithSumOrLess(int[] array, int sum)
	{
	Arrays.sort(array);
	int tripletCount = 0;
	for(int i = 0; i < array.length-2; i++ )
	{
	for(int j = i+1, k = array.length-1; j < k; )
	{
	int curSum = array[i] + array[j] + array[k];
	if( curSum <= sum )
	{
	tripletCount += k-j;
	j++;
	k--;
	}
	else
	{
	k--;
	}
	}
	}
	return tripletCount;
	}
	}

by GitHub