Counting pairs of integers which have a predefined difference

Question

I'm doing this HackerRank problem:

Given \$N\$ integers, count the number of pairs of integers whose difference is \$K\$.

So far, I've ended up with this code:

def countPairs(ar, n, k):
    ar.sort()
    temp = list(ar)
    i = 0
    j = n - 1
    cnt = 0
    while i < n:
        while abs(ar[i] - temp[j]) > k and j > 0:
            j -= 1
        if abs(ar[i] - temp[j]) == k:
            cnt += 1
        i += 1
        j = n - 1  
    return cnt

if __name__ =='__main__':
    N, K = map(int, raw_input().strip().split())
    A = map(int, raw_input().strip().split())
    print countPairs(A, N, K)

The problem is some test cases terminated due to timeout (took more than 10 seconds).

I tried calculating the time complexity. Sorting the list in the first step of countPairs takes \$O(n\log n)\$ and copying the list in the second step into temp takes \$O(n)\$.

This link says that the "two-pointer technique" has \$O(n)\$ complexity. So, the overall complexity is \$O(n \log n)\$. Does anyone know how to make it run faster?

To be honest, I think the "two-pointer technique" isn't \$O(n)\$, but \$O(n^2)\$ because of the two while loops, but I'm not sure.

Answer 1

Style

• 1-letter variable names doesn't have meaning. It is acceptable for i, j or k as loop indexes but other than that it mostly harms readability;
• Same for temp or abreviations (ar, cnt, diff): give these variables some real names;
• functions names should be snake_case instead of camelCase;

• iterating over the values of a sequence is more easily done using a for loop:

  for value in ar:
      # do something with `value` instead of `ar[i]`
  

• Since you’re not modifying it, you can remove temp and access ar[j] directly.

• You don't need to strip() before a split() without arguments, it will be done automatically.
• You correctly used if __name__ == '__main__', that's a good habit to keep.

After that, the code could look like:

def count_pairs(data, length, delta):
    data.sort()
    i = length - 1
    num_pairs = 0
    for value in data:
        while abs(value - data[i]) > delta and i > 0:
            i -= 1
        if abs(value - data[i]) == delta:
            num_pairs += 1
        i = length - 1  
    return num_pairs

if __name__ =='__main__':
    N, K = map(int, raw_input().split())
    data = map(int, raw_input().split())
    print count_pairs(data, N, K)

Complexity

For now, you’re iterating over the same list several times to find the pairs and you need to sort the list beforehand for an efficient computation. But you could compute, for each element, the element that need to be present to form a pair (\$O(n)\$) and then, for each initial element (\$O(n)\$), check if it exists in the array of computed "need-to-be-there" elements.

Luckyly, Python's sets can provide such checks with an \$O(1)\$ performace, for an overall \$O(n)\$ time complexity:

def count pairs(data, delta):
    needed_values = set(value + delta for value in data)
    return sum(value in needed_values for value in data)

Here I abuse the fact that booleans are integers after all, but you could write it in a more proper way using:

def count_pairs(data, delta):
    needed_values = set(value + delta for value in data)
    exist_in_pair = [value for value in data if value in needed_values]
    return len(exist_in_pair)

Answer 2

Okay, so I solved the problem. It seemed like the while loop DID have O(n^2) complexity. I changed the function countPairs() to this:

def countPairs(ar, n, k):
    ar.sort()
    i = 1
    j = 0
    cnt = 0
    while i != n:
        diff = ar[i] - ar[j]
        if diff == k:
            cnt += 1
            i += 1
            j += 1
        elif diff < k:
            i += 1
        else:
            j += 1
    return cnt

Now all the test cases are passed in less than 1 second.