algorithm - 3SUM 问题(寻找三元组)优于 O(n^2)

Question

考虑以下问题:

Given an unsorted array of integers, find all triplets that satisfy x^2 + y^2 = z^2.

For example if given array is 1, 3, 7, 5, 4, 12, 13, the answer should be 5, 12, 13 and 3, 4, 5

我建议使用以下复杂度为 O(n^2) 的算法:

• 将数组降序排列，O(nlogn)
• 对每个元素求平方，O(n)

现在它简化为在排序数组中找到所有三元组 (a,b,c) 的问题，使得 a = b+c。

面试官坚持要一个比 O(n^2) 更好的解决方案。

我已阅读 3SUM problem on Wikipedia ，它强调如果数字在 [-u,u] 范围内，则可以在 O(n+ulogu) 中解决问题，假设数组可以表示为位向量。但是我无法清楚地了解进一步的解释。

有人可以通过一个很好的例子帮助我理解发生了什么吗？

Answer 1

首先。在最坏的情况下找到所有 三元组是O(n^3)。假设您有 n=3k 个数字。其中k个3个，k个4个，k个5个。

3,....,3,4,....,4,5,.....5

有 k^3 = n^3/27 = O(n^3) 这样的三元组。仅打印它们需要 O(n^3) 时间。

接下来将以这种形式解释3SUM问题:

有数字 s_1, ..., s_n，每个数字都在 [-u;u] 范围内。 a+b=c 有多少个三元组 a,b,c？

1. 转变。获取 2*u 个数字 a_-u, ..., a_0, a_1, ... a_u。 a_i 是 s_j 的数量，即 s_j = i。这是在 O(n+u)

中完成的

2. res = a_0 * sum(i=-u..u, i=/=0, C(a_i, 2)) + C(a_0,3) a_0 = 0

3. 构建一个多项式 P(x) = sum(i = -u...u, a_i*x^(i+u)。

4. 使用 FFT 求 Q(x) = P(x)*P(x) .

5. 注意 Q(x) = sum(i=-2u..2u, b_i*x^(i+2*u))，其中 b_i 是 s_u,s_k 对的数量 s_u+s_k = i(这包括使用相同的数字两次)。

6. 对于所有 even i 做 b_i = b_i - a_(i/2)。这将删除使用相同号码两次。

7. 对所有 b_i*a_i/2 求和 - 将其添加到 res。

示例:为了更简单，我假设数字的范围是 [0..u] 并且不会在 x 的幂中使用任何 +u。假设我们有数字 1,2,3 - a_0 = 0, a_1 = 1, a_2 = 1, a_3 = 1

• res = 0

• P(x) = x + x^2 + x^3

• Q(x) = x^2 +2x^3+3x^4+2x^5+x^6

• 减去 b_2 = 0, b_3 = 2, b_4 = 2, b_5 = 2, b_6 = 0

• res += 0*1/2 + 2*1/2 + 2*0/2 + 2*0/2 + 6*0/2 = 1