The question is a cute variation of the 3 sum problem. So lets first go through the classic 2 loop set lookup 3 sum solution.
- We are given a list to get 3 variables (\$a\$, \$b\$ and \$c\$) and a constant (\$d\$). In the original problem \$d = 0\$.
- We know \$a + b + c = d\$, so we can rearrange to find \$c\$. \$c = d - a - b\$
- We then build a set of the input list. If the language doesn't have a set data type a hash map can be abuse instead. As both have an \$O(1)\$ lookup.
- We then just do
input_set.contains(d - a - b) to drop the inner loop.
Here's an example in Python:
# WARNING: the actual code should exclude duplicate indexed values in both the `values_set` lookup and `b`
def is_three_sum(values: list[int], d: int = 0) -> bool:
values_set = set(values)
for a in values:
for b in values:
if d - a - b in values_set:
return True
return False
The variation used is cute as we now need to use modular arithmetic. I have renamed the OP's d to e to not conflate with the \$d\$ in my answer.
- We have the equation \$a + b + c \mod e = d\$. Where \$d = 0\$.
- We can expand to \$(a \mod e) + (b \mod e) + (c \mod e) \mod e = d \mod e\$.
- Rearrange \$(c \mod e) = (d \mod e) - (a \mod e) - (b \mod e) \mod e\$.
- Simplify \$c \mod e = d - a - b \mod e\$.
- Now all we need to do is precompute
values_set to be a set of c % e when we do the lookup.
# WARNING: the actual code should exclude duplicate indexed values in both the `values_set` lookup and `b`
def is_three_sum(values: list[int], d: int = 0, e: int = 1) -> bool:
values_set = {v % e for v in values}
for a in values:
for b in values:
if (d - a - b) % e in values_set:
return True
return False
How to fix the WARNING:
b is simple and is what your code uses.values_set can be built at the end of the first loop to not have duplicate values.
def is_three_sum(values: list[int], d: int = 0, e: int = 1) -> bool:
values_set = set()
for i in range(len(values)):
for j in range(i + 1, len(values)):
if (d - values[i] - values[j]) % e in values_set:
return True
values_set.add(values[i] % e)
return False
NOTE: you can micro-optimise values_set to be an array of size \$e\$ for \$O(1)\$ lookup which avoids hashing.
Finally lets figure out how to change from an 'is' question (bool) to the count question (int).
Let's change return True to count += 1 and see how the code performs with whitebox testing.
For simplicity \$e = 1\$ so all integers sum to 0.
With values = [0, 1, 2]: (expected output: \$\binom{3}{3} = 1\$)
i = 0, values_set = set(), \$j \in \{1, 2\}\$
As values_set is empty none of the inputs can be in the set.
i = 1, values_set = {0}, \$j \in \{2\}\$
j = 2. The moduloed sum of (0, 1, 2) is 0 which is in values_set. count += 1
The output is 1, as expected.
With values = [0, 1, 2, 3]: (expected output: \$\binom{4}{3} = 4\$)
i = 0, values_set = set(), \$j \in \{1, 2, 3\}\$
As values_set is empty none of the inputs can be in the set.
i = 1, values_set = {0}, \$j \in \{2, 3\}\$
j = 2. The moduloed sum of (0, 1, 2) is 0 which is in values_set. count += 1j = 3. The moduloed sum of (0, 1, 3) is 0 which is in values_set. count += 1
i = 2, values_set = {0}, \$j \in \{3\}\$
j = 3. The moduloed sum of (0, 2, 3) is 0 which is in values_set. count += 1
The output is 3, which is not as expected.
The problem is we're missing (1, 2, 3). So when i = 2 and j = 3 we should have done count += 2 to account for both (0, 2, 3) and (1, 2, 3).
Lets change values_set to values_count and change values_set.add(values[i] % e) to values_count[values[i] % e] += 1
With values = [0, 1, 2, 3, 4]: (expected output: \$\binom{5}{3} = 10\$)
i = 0, values_count = [0], \$j \in \{1, 2, 3, 4\}\$
- For all values of
j we do count += 0.
i = 1, values_count = [1], \$j \in \{2, 3, 4\}\$
j = 2. The moduloed sum of (0, 1, 2) is 0 which is 1. count += 1j = 3. The moduloed sum of (0, 1, 3) is 0 which is 1. count += 1j = 4. The moduloed sum of (0, 1, 4) is 0 which is 1. count += 1
i = 2, values_count = [2], \$j \in \{3, 4\}\$
j = 3. The moduloed sum of (0-1, 2, 3) is 0 which is 2. count += 2j = 4. The moduloed sum of (0-1, 2, 4) is 0 which is 2. count += 2
i = 3, values_count = [3], \$j \in \{4\}\$
j = 4. The moduloed sum of (0-2, 3, 4) is 0 which is 3. count += 3
The output is 10, which is expected.
def count_three_sum(values: list[int], d: int = 0, e: int = 1) -> int:
count = 0
values_count = [0] * e
for i in range(len(values)):
for j in range(i + 1, len(values)):
count += values_count[(d - values[i] - values[j]) % e]
values_count[values[i] % e] += 1
return count
NOTE: all code is untested and provided as a reference
I have graphed the performance of the algorithms. To show how going from an \$O(n^3)\$ algorithm to an \$O(n^2)\$ algorithm greatly improves the performance of the code.
def test_orig(a: list[int], d: int = 1):
count = 0
n = len(a)
for i in range(n):
for j in range(i + 1, n):
for k in range(j + 1, n):
if (a[i] + a[j] + a[k]) % d == 0:
count += 1
return count
def test_peil(values: list[int], e: int = 1, d: int = 0) -> int:
count = 0
values_count = [0] * e
for i in range(len(values)):
for j in range(i + 1, len(values)):
count += values_count[(d - values[i] - values[j]) % e]
values_count[values[i] % e] += 1
return count
def test_peil__micro(values: list[int], e: int = 1, d: int = 0) -> int:
count = 0
values_count = [0] * e
values_count[values[0] % e] += 1
for i in range(1, len(values)):
d_ = d - values[i]
for j in range(i + 1, len(values)):
count += values_count[(d_ - values[j]) % e]
values_count[values[i] % e] += 1
return count
import functools
import random
import matplotlib.pyplot
import numpy
import graphtimer
random.seed(42401)
@functools.cache
def args_conv(size: int) -> tuple[list[int], int]:
arr = list(range(10 * int(size)))
random.shuffle(arr)
return arr[:int(size)], random.randrange(1, max(2, int(size) // 10))
def main():
fig, axs = matplotlib.pyplot.subplots()
axs.set_yscale('log')
axs.set_xscale('log')
(
graphtimer.Plotter(graphtimer.MultiTimer([test_orig, test_peil, test_peil__micro]))
.repeat(10, 10, numpy.logspace(0, 2, num=50), args_conv=args_conv)
.min()
).plot(axs, x_label='len(nums)')
fig.show()
input()
if __name__ == '__main__':
main()
