Efficient Inclusion Check

Problem Statement

I recently stumbled upon a cute problem where you have three sets of numbers $$A, B, C \subseteq \lbrace 0, \dots, n \rbrace.$$ You want to answer if there exists $$a \in A, b \in B \text{ such that } c = a + b \in C.$$ It's a conceptually simple problem: are there two numbers from the first two sets that sum to a number in the third set? The hard part is we want to solve in it O(n log n) time.

Brute force search

The simplest solution is to define the C set as a set in Python so that we can do constant time inclusion checking.

for a in A:
  for b in B:
    if a+b in C:
      return True
return False

Thus the above code is O(n²), not our goal yet. (Or if we don't have efficient checking of inclusion in C, it's O(n³).)

Clever search

We can do better! The observation that we can represent the sets as polynomials is critical. Thus let, $$A = \lbrace a_1, a_2, a_3, \dots, a_n \rbrace$$ be represented as $$p_A(x) = x^{a_1} + x^{a_2} + \dots + x^{a_n}.$$ Do the same for B and C. Then, consider $$p_A(x) \times p_B(x).$$

Let's run a little example to see why this matters. $$ A = \lbrace 1, 4, 9 \rbrace, B = \lbrace 2, 10 \rbrace, C = \lbrace 12, 1000 \rbrace$$ Then, $$p_A(x) = x + x^4 + x^9$$ and $$p_B(x) = x^2 + x^{10}$$ and $$p_C(x) = x^{12} + x^{1000}.$$ Notice though that $$p_A(x) \times p_B(x) = x^3 + x^6 + 2 x^{11} + x^{14} + x^{19}.$$ We only have cofficients greater than zero for terms that are the sums of elements in A and B! Thus, if we can efficiently calculate the polynomial product, we can just compare the coefficient list to C.

Efficient polynomial multiplication

It turns out multiplying polynomials in O(n log n) time is possible if we do it in Fourier space.

from numpy.fft import fft, ifft
from numpy import real, imag

def polynomial_multiply(a_coeff_list, b_coeff_list):
  """
  Return the coefficient list of the multiplication
  of the two polynomials
  """
    n, m = len(a_coeff_list), len(b_coeff_list)

    # we pad with zeros since c is size m + n - 2
    a_fft_coeffs = fft(a_coeff_list + [0.0]*(m - 1))
    b_fft_coeffs = fft(b_coeff_list + [0.0]*(n - 1))

    c_fft_coeffs = [a*b for a, b in zip(a_fft_coeffs, b_fft_coeffs)]
    c = [real(v) for v in ifft(c_fft_coeffs)]
    return c

The final solution

def check_sum_exists(a, b, c, n):
  """
  Inputs:
  sets a, b, c
  n is the maximum number in a, b, c together
  """
  # set up the polynomial coefficients
  a_coeffs = [0]*n
  b_coeffs = [0]*n
  for aa in a:
    a_coeffs[aa] = 1
  for bb in b:
    b_coeffs[bb] = 1

  # multiply them together
  c_coeffs = polynomial_multiply(a_coeffs, b_coeffs)

  # Make sure the coefficients are integers without floating point errors
  coeffs_copy = []
  for num in c_coeffs:
    if(abs(num-0) < abs(num-1)):
      coeffs_copy.append(0)
    elif(abs(num-1) < abs(num-2)):
      coeffs_copy.append(1)
    else:
      coeffs_copy.append(2)

  return any([coeffs_copy[cc] >= 1 for cc in c])

Conclusion

This cute solution is an approach that I wouldn't arrived at initially. Just because you don't see a more efficient solution doesn't mean there isn't one.