python - 寻找零和的三元组

Question

我正在尝试解决 GeeksClasses 中的问题，但我的提交有问题。我的代码有效，但他们说您的程序花费的时间比预期的要长。

问题链接: https://practice.geeksforgeeks.org/problems/find-triplets-with-zero-sum/1/?track=SPCF-Sorting&batchId=154

问题陈述:

给定一个整数数组。检查它是否包含总和为零的三元组。

输入:

输入的第一行包含一个整数T，表示测试用例的数量。然后是 T 个测试用例。每个测试用例的第一行包含一个整数 N，表示数组中元素的数量。每个测试用例的第二行包含数组的 N 个空格分隔值。

输出

对于每个测试用例，如果存在三元组，则输出为 1，否则为 0

预期辅助空间:O(1)

预期时间复杂度:O(n2)

示例:

输入:

2

5

0 -1 2 -3 1

3

1 2 3

输出:

1

0

这是我的代码

def isPair(arr,left,right,u):
    while left < right:
        if arr[left] + arr[right] < u:
            left += 1
        elif arr[left] + arr[right] == u:
            return True
        elif arr[left] + arr[right] > u:
            right -= 1
    return False

def findTriplets(a,n):
    #code here
    a = sorted(a)
    for i in range(n):
        if isPair(a,i+1,n-1,0-a[i]):
            return 1
    return 0
#driver code
if __name__=='__main__':
    t=int(input())
    for i in range(t):
        n=int(input())
        a=list(map(int,input().strip().split()))
        print(findTriplets(a,n))

Answer 1

这个问题看起来很有趣，这里有两个我们可以使用的观察结果。每个有效的三元组都是以下任一形式:

1. (0, -x, x)
2. 或 (x, y, z) 使得 x 和 y 与 z 的符号相反并且 x + y = - z

我会考虑一种更简单的输入形式，因为您的大部分输入对于两个整数列表的内容都是多余的，即。 example_1 = [[0, -1, 2, -3, 1], [1, 2, 3]]应该导致 [1, 0] .

鉴于我认为以下是一个相当快速/可读的解决方案:

from itertools import combinations

def solve_all(inputs):
    return [solve(i) for i in inputs]

def solve(single_input):
    input_set = set(single_input)
    negatives_set = set(-x for x in single_input if x < 0)
    positives_set = set(x for x in single_input if x > 0)

    if 0 in input_set and len(negatives_set & positives_set) > 0:
        return 1

    if any(sum(c) in positives_set for c in combinations(negatives_set, 2)):
        return 1

    if any(sum(c) in negatives_set for c in combinations(positives_set, 2)):
        return 1

    return 0