3 sum implementation in python 2.7

Question

Working on below problem, looking for advice on code style, functional bug and time complexity improvements.

Also looking for advice in,

1. If array is not sorted, any solutions available for O(n^2) time complexity?
2. In my loop, I always search from exclude_index + 1 to reduce duplicate, is it a safe start boundary?

Problem,

Given a sorted array and a target number, tell whether there are three numbers in the array which add up to the target number. (Time complexity should be O(n^2))

Source code,

def two_sum(numbers, start, end, target, exclude_index):
 while start < end:
 if start == exclude_index:
 start += 1
 elif end == exclude_index:
 end -= 1
 elif numbers[start] + numbers[end] > target:
 end -= 1
 elif numbers[start] + numbers[end] == target:
 print numbers[start], numbers[end], numbers[exclude_index]
 start += 1
 end -= 1
 while start < end and (numbers[start] == numbers[start - 1] or start == exclude_index):
 start += 1
 while start < end and (numbers[end] == numbers[end + 1] or end == exclude_index):
 end -= 1
 else:
 start += 1
def three_sum(numbers, target):
 for exclude_index in range(len(numbers)):
 if exclude_index > 0 and numbers[exclude_index] == numbers[exclude_index - 1]:
 continue
 two_sum(numbers, exclude_index + 1, len(numbers) - 1, target - numbers[exclude_index], exclude_index)
if __name__ == "__main__":
 numbers = [1, 3, 4, 6, 7, 8, 9, 9, 10, 12]
 three_sum(numbers, 15)

Answer 1

As with most of your questions, your code violates PEP8.

---

1. Don't print your output.
2. Since \$\text{start} = \text{exclude_index} + 1\$ you can never visit exclude_index.

Your code with the above changes make your code more readable. But it'd be better if you:

1. Start using Python's standard library. Itertools can drastically improve your codes readability.
2. Stop relying on while so much.
3. Use a dictionary, it's fast, and increases readability.
4. sorted runs in \$O(n\log(n))\$ time.

Using itertools.combinations_with_replacement and collections.Counter you can generate all three of your numbers, still in \$O(n^2)\$ time.

1. Pass a sorted list to combinations_with_replacement, and generate two numbers.
2. Check target - i - j is in the list of numbers, and that it is greater or equal to i and j.
3. Group; i, j and k, in a Counter and check there is enough of each number in the original list.

import collections
import itertools
def three_sum(numbers, target):
 Counter = collections.Counter
 count = Counter(numbers)
 for i, j in itertools.combinations_with_replacement(list(sorted(count)), 2):
 k = target - i - j
 if k in count and j <= k >= i:
 nums = (i, j, k)
 if all(count[k] >= v for k, v in Counter(nums).items()):
 yield nums

• Lin Ma: Thanks Peilonrayz, appreciate you pointed out so many areas for improvements and vote up. There is only one thing I do not quite agree, which is if you enumerate all 3 number combinations, time complexity is O(n^2) which is less than mine (which is O(n*lgn))?
• Lin Ma: BTW, another thing I am confused is, why you are using combinations_with_replacement other than itertools.combinations, I think combinations_with_replacement considers repeated elements, but list(sorted(count)) should be unique which is dict keys?
• Peilonrayz: @LinMa Both mine and yours are \$O(n^2)\$. Try the code without combinations_with_replacement, and you'll see it doesn't work in all cases, such as ones that have duplicate values.