Friday, December 23, 2011

Finding all unique triplets that sums to zero

Source

Given a set S of n integers, are there elements a, b, c in S such that a + b + c = 0? Find all unique triplets in the set which gives the sum of zero.
For example, given set S = {-1 0 1 2 -1 -4},
One possible solution set is:
(-1, 0, 1)
(-1, 2, -1)
Note that (0, 1, -1) is not part of the solution above, because (0, 1, -1) is the duplicate of (-1, 0, 1). Same with (-1, -1, 2).
For a set S, there is probably no “the” solution, another solution could be:
(0, 1, -1)
(2, -1, -1)
This is also known as the 3sum problem. The 3sum problem is the extension of the problem below:
Given a set S of n integers, find all pairs of integers of a and b in S such that a + b = k?
The above problem can be solved in O(n) time, assuming that the set S is already sorted. By using two index first and last, each pointing to the first and last element, we look at the element pointed by first, which we call A. We know that we need to find B = k – A, the complement of A. If the element pointed by last is less than B, we know that the choice is to increment pointer first by one step. Similarly, if the element pointed by last is greater than B, we decrement pointer last by one step. We are progressively refining the sum step by step. Since each step we move a pointer one step, there are at most n steps, which gives the complexity of O(n).
By incorporating the solution above, we can solve the 3sum problem in O(n^2) time, which is a straight forward extension.
set<vector<int> > find_triplets(vector<int> arr) {
  sort(arr.begin(), arr.end());
  set<vector<int> > triplets;
  vector<int> triplet(3);
  int n = arr.size();
  for (int i = 0;i < n; i++) {
    int j = i + 1;
    int k = n - 1;
    while (j < k) {
      int sum_two = arr[i] + arr[j];
      if (sum_two + arr[k] < 0) {
        j++;
      } else if (sum_two + arr[k] > 0) {
        k--;
      } else {
        triplet[0] = arr[i];
        triplet[1] = arr[j];
        triplet[2] = arr[k];
        triplets.insert(triplet);
        j++;
        k--;
      }
    }
  }
  return triplets;
}
Note that a set is chosen to store the triplets, because we are only interested in unique triplets. Since the set S is already sorted, and we don’t look back as it progresses forward, we can guarantee there will be no duplicate triplets (Even though the set might have duplicate elements.)

No comments:

Post a Comment