#!/usr/bin/python
# -*- coding: utf-8 -*-
"""Riffing on
<http://nathanleclaire.com/blog/2013/10/22/three-elements-that-sum-to-zero/>.
Given an array of values, find the *array indices* of three of them
that sum to zero.

"""

# This version fits in a Tweet with a few tricks, but this way it's a
# little clearer.
def simple(x):
    return ((i, j, k)
            for i in range(len(x))
            for j in range(i+1, len(x))
            for k in range(j+1, len(x))
            if x[i] + x[j] + x[k] == 0
            ).next()

def test(f):
    assert f([-1, 6, 8, 9, 10, -100, 78, 0, 1]) == (0, 7, 8)

test(simple)

# A version of Nathan's "suave" algorithm using dicts instead of
# binary search (thus theoretically reducing the big-O factor from
# logarithmic to constant, but whatevs.)  Note that this version has
# the unfortunate property that it will return (0, 1, 1) if given
# [6,-3] as input.
def suave(x):
    xi = dict((xv, i) for i, xv in enumerate(x))
    return ((i, j, xi[-x[i]-x[j]])
            for i in range(len(x))
            for j in range(i+1, len(x))
            if -x[i]-x[j] in xi
            ).next()

test(suave)

# Tweetable suave. Almost:
def ts0(x):y=dict((j,i)for i,j in enumerate(x));return((i,j,y[-x[i]-x[j]])for i in range(len(x))for j in range(i+1,len(x))if-x[i]-x[j]in y).next()

# Actually tweetable, and even using Python lists, it's still O(N²),
# because although it does the extra work of making a bunch of partial
# copies of the list, it only makes N of them.  But that inefficiency
# is still sort of offensive, particularly since it's only there to be
# able to use enumerate().  But if you apply it to a Numpy array
# instead, that inefficiency vanishes.
def ts(x):e=enumerate;y=dict((j,i)for i,j in e(x));return((i,i+j+1,y[-xi-xj])for i,xi in e(x)for j,xj in e(x[i+1:])if-xi-xj in y).next()

test(ts)