LeetCode 532. K-diff Pairs in an Array

September 13, 2020

We can do this in $O(n)$ time and $O(n)$ space with a hash table:

For each number, $x$ , check if $(x-k)$ or $(x+k)$ is in our hash table.
Add $x$ to our hash table.

Hand-wavey argument for correctness:

$(x-k)$ and $(x+k)$ are the only possible pairs for $x$ and we check both.
Might we be missing pairs? Suppose that, for $x$ , we aren't checking a possible pair. That missed element can't be from earlier in the list since those elements are all in our hash table which we're checking against. So the missed element must be after $x$ in the list. However, once our iteration reaches that missed element, it would find $x$ in the hash table. By contradiction, there are no pairs that we might miss.

We can't do better than $O(n)$ since we have to look at each element at least once if we're finding all pairs.

from typing import List
import unittest

class Test(unittest.TestCase):

    def setUp(self):
        self.s = Solution()

    def test_examples(self):
        self.assertEqual(self.s.findPairs([3, 1, 4, 1, 5], 2), 2)
        self.assertEqual(self.s.findPairs([1, 2, 3, 4, 5], 1), 4)
        self.assertEqual(self.s.findPairs([1, 3, 1, 5, 4], 0), 1)

    def test_empty(self):
        self.assertEqual(self.s.findPairs([], 0), 0)

    def test_negative_k(self):
        self.assertEqual(self.s.findPairs([1, 2], -1), 0)

    def test_repeats(self):
        self.assertEqual(self.s.findPairs([1, 1, 1, 2, 1, 1], 1), 1)

    def test_abs_diff(self):
        self.assertEqual(self.s.findPairs([2,1,0,5], 1), 2)

class Solution:

    def findPairs(self, nums: List[int], k: int) -> int:
        if k < 0:
            return 0
        pairs = set() # store pairs as {(a,b), a <= b} so they're unique
        seen = set()
        for n in nums:
            if n-k in seen:
                pairs.add((n-k, n))
            if n+k in seen:
                pairs.add((n, n+k))
            seen.add(n)
        return len(pairs)

if __name__ == "__main__":
    unittest.main()