ROSALIND｜3SUM (3SUM)

發表於

給定 k 個長度為 n 的整數數列，求各數列裡總和為零的三個數字的位置。若存在多組解，則以自位置順序以最靠前的為準。

Given: A positive integer $k \leq 20$, a positive integer $n \leq 10^4$, and $k$ arrays of size $n$ containing integers from $-10^5$ to $10^5$.

Return: For each array $A[1..n]$, output three different indices $1 \leq p < q < r \leq n$ such that $A[p] + A[q] + A[r] = 0$ if they exist, and "-1" otherwise.

(https://rosalind.info/problems/3sum/)

LeetCode 3sum 的要求和 Rosalind 不同，這題只需要回傳第一組解即可，不用顧慮重複解的狀況。

3sum 問題可以化約為一系列 2sum 問題。我們首先考慮數列裡其中一個數字 l[i]，3sum 的目標是尋找另外兩個數字 l[j] 和 l[k] 使 l[i] + l[j] + l[k] = target。

這個關係式可以改寫為 l[j] + l[k] = target - l[i]。由於等號右側的資訊已知，能使用 2sum 的解法判斷 l[j] 和 l[k] 是否存在。依據檢查數列中每個數字與 target 的差值是否存在 2sum 解，即可得出 3sum 的解。

為了避免正在檢查的數字干擾計算，我們要在 2sum 加入新的參數 skip_index，表示這個數字不在搜尋範圍內。

1
2
3
4
5
6
7
8
9
10
11
def two_sum(l, target, skip_index=None):
    diffs = {}
    for i, num in enumerate(l):
        # 3sum 的候選數字不納入 2sum 計算。
        if i == skip_index:
            continue
        diff = target - num
        if diff in diffs:
            return (diffs[diff], i)
        diffs[num] = i
    return None

因為 Rosalind 這題僅要求第一組解，所以找到答案時直接回傳即可。

1
2
3
4
5
6
7
def three_sum(l, target):
    for i, num in enumerate(l):
        residual = target - num
        out = two_sum(l, residual, skip_index=i)
        if out is not None:
            return (i, out[0], out[1])
    return None