算法题/双指针：和为0的三数

题目

给定一个未排序数数组，找出其中所有相加为0的不同的三元组。

示例1：

输入：[-3, 0, 1, 2, -1, 1, -2]
输出：[-3, 1, 2], [-2, 0, 2], [-2, 1, 1], [-1, 0, 1]
解释：上面4个集合的数的和为0

示例2：

输入：[-5, 2, -1, -2, 3]
输出：[-5, 2, 3], [-2, -1, 3]
解释：上面2个集合的数的和为0

解答1：

Java

import java.util.*;

class TripletSumToZero {

  public static List<List<Integer>> searchTriplets(int[] arr) {
    Arrays.sort(arr);
    List<List<Integer>> triplets = new ArrayList<>();
    for (int i = 0; i < arr.length - 2; i++) {
      if (i > 0 && arr[i] == arr[i - 1]) //跳过相同的元素去避免重复的三元组
        continue;
      searchPair(arr, -arr[i], i + 1, triplets);
    }

    return triplets;
  }

  private static void searchPair(int[] arr, int targetSum, int left, List<List<Integer>> triplets) {
    int right = arr.length - 1;
    while (left < right) {
      int currentSum = arr[left] + arr[right];
      if (currentSum == targetSum) { //找到三元组
        triplets.add(Arrays.asList(-targetSum, arr[left], arr[right]));
        left++;
        right--;
        //此处需要left<right是因为下面需要left++，而left++之前一定要检查left是否小于right
        while (left < right && arr[left] == arr[left - 1])
          left++; //跳过相同的元素避免相同的三元组
        while (left < right && arr[right] == arr[right + 1])
          right--; //跳过相同的元素避免相同的三元组
      } else if (targetSum > currentSum)
        left++; //我们将left指针右移来得到一个更大的数，从而得到更大的currentSum
      else
        right--; //我们将right指针左移来得到一个更小的数，从而得到更小的currentSum
    }
  }

  public static void main(String[] args) {
    System.out.println(TripletSumToZero.searchTriplets(new int[] { -3, 0, 1, 2, -1, 1, -2 }));
    System.out.println(TripletSumToZero.searchTriplets(new int[] { -5, 2, -1, -2, 3 }));
  }
}

JavaScript

function searchTriplets(arr) {
  arr.sort((a, b) => a - b);
  const triplets = [];
  for (let i = 0; i < arr.length - 2; i++) {
    if (i > 0 && arr[i] === arr[i - 1])
      continue;
    searchPair(arr, -arr[i], i + 1, triplets);
  }
  return triplets;
}

function searchPair(arr, targetSum, left, triplets) {
  let right = arr.length - 1;
  while (left < right) {
    const currentSum = arr[left] + arr[right];
    if (currentSum === targetSum) {
      triplets.push([-targetSum, arr[left], arr[right]]);
      left++;
      right--;
      while (left < right && arr[left] === arr[left - 1])
        left++;
      while (left < right && arr[right] === arr[right + 1])
        right--;
    } else if (targetSum > currentSum) {
      left++;
    } else {
      right--;
    }
  }
}

console.log(searchTriplets([-3, 0, 1, 2, -1, 1, -2]));
console.log(searchTriplets([-5, 2, -1, -2, 3]));

Python

def searchTriplets(arr):
    arr.sort()
    triplets = []
    for i in range(len(arr) - 2):
        if i > 0 and arr[i] == arr[i - 1]:
            continue
        searchPair(arr, -arr[i], i + 1, triplets)
    return triplets

def searchPair(arr, targetSum, left, triplets):
    right = len(arr) - 1
    while left < right:
        currentSum = arr[left] + arr[right]
        if currentSum == targetSum:
            triplets.append([-targetSum, arr[left], arr[right]])
            left += 1
            right -= 1
            while left < right and arr[left] == arr[left - 1]:
                left += 1
            while left < right and arr[right] == arr[right + 1]:
                right -= 1
        elif targetSum > currentSum:
            left += 1
        else:
            right -= 1

print(searchTriplets([-3, 0, 1, 2, -1, 1, -2]))
print(searchTriplets([-5, 2, -1, -2, 3]))

Go

package main

import (
	"fmt"
	"sort"
)

func searchTriplets(arr []int) [][]int {
	sort.Ints(arr)
	triplets := [][]int{}
	for i := 0; i < len(arr)-2; i++ {
		if i > 0 && arr[i] == arr[i-1] {
			continue
		}
		searchPair(arr, -arr[i], i+1, &triplets)
	}
	return triplets
}

func searchPair(arr []int, targetSum int, left int, triplets *[][]int) {
	right := len(arr) - 1
	for left < right {
		currentSum := arr[left] + arr[right]
		if currentSum == targetSum {
			*triplets = append(*triplets, []int{-targetSum, arr[left], arr[right]})
			left++
			right--
			for left < right && arr[left] == arr[left-1] {
				left++
			}
			for left < right && arr[right] == arr[right+1] {
				right--
			}
		} else if targetSum > currentSum {
			left++
		} else {
			right--
		}
	}
}

func main() {
	fmt.Println(searchTriplets([]int{-3, 0, 1, 2, -1, 1, -2}))
	fmt.Println(searchTriplets([]int{-5, 2, -1, -2, 3}))
}