#!/usr/bin/python
# -*- coding: utf-8 -*-
"""Riffing on
.
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)