Find all combinations of length 3 whose sum is divisible by a given number

Question

I came up with a suitable solution to a HackerRank problem that failed because the execution took longer than 10 seconds on lists that were of very large size.

The problem: Given a list arr and an integer d, find all combinations of length 3 whose sum is divisible by d

My solution:

[
    sum(p) % d
    for p in combinations(arr, r=3)
].count(0)

I simply cannot find a way to optimize this (for time) any more. Can this be further optimized for speed?

Answer 1

Time complexity of iterating through all the combinations is O(n³).

We can solve it in O(n²) by using the fact that if you know 2 numbers out of 3, you know exactly what modulo-d remainder you need from the 3rd number. If the tests for the problem are properly written, none of the O(n³) solutions will pass, so the only way to get the solution accepted is to go for a better time complexity. So here is how it works.

Put modulo-d remainder of every element in a hashset. Iterate through all possible pairs with a double for loop and for each of them check whether d - (a + b) % d is in the hashset. Such lookups only take O(1) which makes your overall complexity O(n²) (for the double for loop). If the hashset contains the number you're looking for, you have found a combination.

Answer 2

As far as clean code which expresses the intent goes, I quite like your little snippet. In fact, it's the sort of thing which, if I were solving this problem in production code, I'd probably use your suggestion as a reference implementation. It is as near as possible a direct translation of the problem text into Python code. As such its great value is that it's obviously correct, in the sense of the Tony Hoare maxim that you can either go for simplicity with obviously no errors, or complexity with no obvious errors. That makes it valuable, even if it is too slow for production code, for use in tests to validate production code against.

In terms of speeding it up, a few thoughts come to mind immediately.

1. Minor: Constructing a whole list only to collapse it later may be less efficient than keeping count in a running tally.
2. Major: combinations(arr, r=3) is a deceptively short way of writing three nested for loops, which means the overall complexity is cubic in the length of arr. That's where you're really losing your time, so you'll have to avoid actually working your way through all the possible combinations one by one. I don't want to ruin the exercise by providing a complete solution, but two (related) ways to get that number down would be to explore dynamic programming to reuse work from previous partial combinations or equivalence classes and combinatorics to count functionally similar elements of arr all at once.