如何实现低于O(N²)时间复杂度的过原点2D点对计数？

问题回顾 #

给定N个二维坐标点，统计无序点对(i,j)（i≠j）的数量，使得两点连线经过原点(0,0)。

核心思路 #

两点连线过原点的本质是：两点位于同一条过原点的直线上。我们可以通过归一化点的方向向量，将同一直线上的点归为一组，再用组合数计算每组内的有效点对数量，同时单独处理原点的特殊情况。

关键数学推导 #

两点(x₁,y₁)和(x₂,y₂)连线过原点的等价条件是：x₁*y₂ = x₂*y₁（交叉相乘避免除法精度问题）。

对于非原点的点，我们可以将其坐标归一化为最简方向向量：

1. 计算坐标(x,y)的绝对值的最大公约数gcd(|x|, |y|)
2. 将x和y分别除以gcd，得到最简形式(rx, ry)
3. 统一符号：确保rx为正（若rx=0则ry为正），避免同一直线的点因符号不同被归为不同组（比如(-2,-4)和(1,2)应属于同一组）

优化代码实现 #

import math
from collections import defaultdict

def find_pairs(array, size):
    direction_counts = defaultdict(int)
    origin_count = 0
    
    for x, y in array:
        if x == 0 and y == 0:
            origin_count += 1
            continue
        # 计算最大公约数，注意取绝对值
        g = math.gcd(abs(x), abs(y))
        rx = x // g
        ry = y // g
        # 统一符号：确保rx为正，rx为0时ry为正
        if rx < 0 or (rx == 0 and ry < 0):
            rx = -rx
            ry = -ry
        direction_counts[(rx, ry)] += 1
    
    total = 0
    # 计算非原点组内的点对数量：C(k,2) = k*(k-1)/2
    for cnt in direction_counts.values():
        total += cnt * (cnt - 1) // 2
    # 计算原点相关的点对：原点与所有非原点的点对 + 原点之间的点对
    total += origin_count * (size - origin_count)
    total += origin_count * (origin_count - 1) // 2
    
    return total

复杂度分析 #

• 时间复杂度：O(N log M)，其中M是点坐标的最大值。遍历所有点是O(N)，每个点计算gcd的时间是O(log M)，整体远低于O(N²)。
• 空间复杂度：O(K)，K是不同方向向量的数量，最坏情况下是O(N)（所有点方向都不同），但实际远小于暴力解法的内存占用。

原代码的问题分析 #

1. 暴力解法O(N²)超时：当N较大（如1e4以上）时，双重循环的计算量会急剧增长，无法通过时间限制测试。
2. itertools.combinations内存溢出：生成所有点对需要O(N²)的内存空间，当N较大时会直接超出内存限制。
3. 数学条件冗余：原代码的条件x1*(y2-y1) == y1*(x2-x1)等价于x1*y2 == x2*y1，后者更简洁高效。

内容的提问来源于stack exchange，提问作者user662650