Given two arrays `a` and `b`, find all pairs `(i, j)`such that `i <= j` and `a[i] - b[j] = a[j] - b[i]`

NOTE: Yes, this is for a coding challenge, which I already finished, but I want to see if my answer was correct, and if there are more elegant solutions

I got this question on a coding assignment, and I failed to find the optimal O(n) solution

Given two arrays a and b, find all pairs (i, j) such that i <= j and a[i] - b[j] = a[j] - b[i]. a and b will be the same length at all times.

Give an O(n) solution

The obvious solution is the double for loop:

def simple(a, b):
    count = 0
    for i in range(len(a)):
        for j in range(len(b)):
            if (a[i] - b[j] == a[j] - b[i] and i <= j):
                count += 1
    return count
a = [2, -2, 5, 3]
b = [1, 5, -1, 1]
print('Pairs:', simple(a, b)) # This prints 6, the pairs being [[0, 0], [0, 1], [1, 1], [2, 2], [2, 3], [3, 3]]

But this solution is O(n^2)

Here is my O(n) solution:

I noticed that a[i] - b[j] = a[j] - b[i] can be written as a[i] + b[i] = a[j] + b[j]

So i decided to create two lists list_i and list_j, both the same length as a and b.

In list_i I stored all the results of a[i] + b[i]

In list_j I stored all the results of a[j] + b[j]

Then all I had to do was compare the values at each index for the 2 lists mentioned above and make sure they were the same and that i <= j. The code is below, but it's incorrect.

def complex(a, b):
    n = len(a)
    count = 0
    list_i = [0] * n
    list_j = [0] * n

    for i in range(n):
        list_i[i] = a[i] + b[i]

    for j in range(n):
        list_j[j] = a[j] + b[j]

    i, j = 0, 0
    while (i <= j and i < n and j < n):
        if (list_i[i] == list_j[j]):
            count += 1
            j += 1
        else:
            i = j
    return count

While doing that, I noticed that the index where i and j are the same, will always be true.

So I don't have to check them. I can simply add the line of a to the final count.

I also noticed that I only needed to check adjacent values such that i = i + 1.

This meant that I only had to have one list of sums and in a for loop check to see if sum_list[i] == sum_list[i + 1]

So my final code is as follows:

def complex(a, b):
    n = len(a)
    count = 0
    list_sum = [0] * n

    for i in range(n):
        list_sum[i] = a[i] + b[i]

    for i in range(n - 1):
        if list_sum[i] == list_sum[i + 1]:
            count += 1
    return count + len(a)
a = [2, -2, 5, 3]
b = [1, 5, -1, 1]
print('Pairs:', complex(a, b)) # This prints 6, the pairs being [[0, 0], [0, 1], [1, 1], [2, 2], [2, 3], [3, 3]]

This gives the correct solution and is O(n), but is it correct?

Also is there a better solution using sets or dicts?

If you're only interested in the number of pairs, not the pairs themselves, this can be solved in O(n) with a dictionary.

First, let c be the elementwise sum of a and b (i.e. c = [ai + bi for ai, bi in zip(a, b)]). Observe that the constraint a[i] + b[i] == a[j] + b[j] implies that we're looking for i, j such that c[i] == c[j].

For each unique value v of c, if there are n indices such that c[i] == v, then there will be n * (n + 1) // 2 pairs of indices such that c[i] == c[j] and i <= j. Therefore, all we need to do is count the number of indices per unique value of c and apply this formula:

from collections import Counter

def count_pairs(a, b):
    c = [ai + bi for ai, bi in zip(a, b)]
    ctr = Counter(c)
    return sum(n * (n + 1) // 2 for n in ctr.values())

This runs in O(n) amortized time: constructing c takes O(n) time, constructing ctr takes O(n) amortized time (it is a dict which stores only the count of each unique value), and computing the sum takes O(n) time (the calculation runs once per unique value in c, and there are at most n such unique values).

An easy O(N) solution can be found by using a hashmap, which allows O(1) lookup.

> I also noticed that I only needed to check adjacent values

This is true if the list_sum list is sorted, but it's not true in the general case (e.g. if list_sum is [10, 20, 10], coming from e.g. a = [1, 2, 3], b = [9, 18, 7]).

So to fix your code, do a list_sum.sort() before your 3rd for loop. This will make it slower though, because the sorting takes O(N * log(N)) time.

Also you may have to count longer runs of the same value more. Your complex function returns 5 if list_sum is [10, 10, 10], but shouldn't it return 6 instead?

"Give an O(n) solution"

Hmmmm... consider cases where a and b are both arrays containing n 0s each, e.g.:

n = 2

a = [ 0, 0 ]
b = [ 0, 0 ]
solution = [ [ 0, 0 ], [ 0, 1 ], [ 1, 1 ] ]
solution length: 3

n = 3

a = [ 0, 0, 0 ]
b = [ 0, 0, 0 ]
solution = [ [ 0, 0 ], [ 0, 1 ], [ 0, 2 ], [ 1, 1 ], [ 1, 2 ], [ 2, 2 ] ]
solution length: 6

n = 4

a = [ 0, 0, 0, 0 ]
b = [ 0, 0, 0, 0 ]
solution = [ [ 0, 0 ], [ 0, 1 ], [ 0, 2 ], [ 0, 3 ], [ 1, 1 ], [ 1, 2 ], [ 1, 3 ], [ 2, 2 ], [ 2, 3 ], [ 3, 3 ] ]
solution length: 10

In general when a and b are arrays of n 0s, the length of the solution is the triangular numbers; f(n) = n * (n + 1) * 1/2 - this is quadratic growth, and doesn't look to me like it will ever be captured in less than O(n^2) time.

Are we missing some details from the question?

• e.g. return the number of pairs, not the pairs themselves?
• are a and b sorted?
• are items in a and b guaranteed to be unique?