python - 选择数字对的最快算法

Question

问题: 给定 N 个整数 [N<=10^5]，计算差值为 K 的整数对总数。[K>0 且 K<1e9]。 N 个整数中的每一个都将大于 0，并且至少偏离 2^31-1 K(一切都可以用 32 位整数来完成)。

第一行包含 N 和 K(整数)。 第二行包含集合的 N 个数字。确保所有 N 个数字都是不同的。

现在问题来自hackerrank。我得到了问题的解决方案，但它不满足所有示例测试用例的时间限制。我不确定是否可以使用另一种算法，但我没有主意。如果有人花一点时间检查我的代码并提供一两个提示，我将不胜感激。

temp = input()
temp = temp.split(" ")
N = int(temp[0])
K = int(temp[1])
num_array = input()
num_array = num_array.split(" ")
diff = 0
pairs= 0
i = 0
while(i < N):
    num_array[i] = int(num_array[i])
    i += 1
while(num_array != []):    
    j = 0
    while(j < (len(num_array)-1)):
        diff = abs(num_array[j+1] - num_array[0])
        if(diff == K):
            pairs += 1
        j += 1
    del num_array[0]
    if(len(num_array) == 1):
        break
print(pairs)

Answer 1

您可以按照以下过程在近似线性时间内完成此操作:

所以，O(n) 解:

1. 对于每个数字 x 将其添加到哈希集 H[x]
2. 对于每个数字 x 检查 x-k 是否在 H 中，如果是 - 加 1 回答

或者通过在 O(nlgn) 中使用一些平衡结构(如基于树的集合)

此解决方案基于整数不同 的假设，如果它们不是，您需要存储元素已“添加到集合”的次数，而不是将 1 加到答案中 - 添加H[x]*H[x+k]的乘积

所以通常你会使用一些“默认值为 0”的 HashMap H

1. 对于每个数字x更新映射:H[x] = H[x]+1
2. 对每个数字 x 添加答案 H[x]*H[x-k](你不必检查它是否在 map 中，因为如果不在，则 H[x-k]=0 )

再次 - 使用散列图的解决方案是 O(n)，使用 TreeMap 的 O(nlgn)

给定一组数字 A 和数字 k(不同数字的解决方案):

H=set()
ans=0
for a in A: 
  H.add(a)
for a in A: 
  if a-k in H: 
    ans+=1
print ans

或更短

H=set(A)
ans = sum(1 for a in A if a-k in H)
print ans