3sum.py

from random import SystemRandom
r=SystemRandom()

"""
TASK.
An array A of integers is given. Find all the combinations of indexes i,j,k, so A[i]+A[j]+A[k]==0. Preferably in O(N^2).
It's a variant of 3SUM puzzle.
"""

def bad_algorithm(INPUT):
	"""
	O(N^3)
	"""
	results = set()

	for i,I in enumerate(INPUT):
		for j,J in enumerate(INPUT):
			for k,K in enumerate(INPUT):
				if i == j or i == k or j == k:
					pass
				else:
					if I+J+K == 0:
						results.add(tuple(sorted([i,j,k])))

	return results

def quicksort_memory(inpoot):
	"""
	Sorts a list and also returns a list of original indicies of values.
	Meaning that store[0] is the index of an element that is located at 0 in sorted list that it had in the unsorted list.
	"""
	store = list(range(len(inpoot)))
	inp = list(inpoot)
	def _quicksort(start, end):
		def swap(i,j):
			inp[i], inp[j] = inp[j], inp[i]
			store[i], store[j] = store[j], store[i]


		if end-start <= 0:
			return inp

		P = start
		L = start + 1
		R = end

		while(L <= R):
			while L <= end and inp[L] < inp[P]:
				L+=1
			if inp[R] >= inp[P]:
				R-=1
			else:
				if L<=R:
					swap(L,R)

		swap(R,P)

		_quicksort(start, R-1)
		_quicksort(R+1, end)

	_quicksort(0, len(inp)-1)

	return inp, store


def better_algorithm(INPUT):
	"""
	O(N^2). Probably...
	"""
	results = set()

	inp, origs = quicksort_memory(INPUT) # array must be sorted
	N = len(inp)

	#till 3rd from end inclusive
	#let's fix the first number and then iterate over other two
	for i in range(0, N-2):
		a = inp[i]

		if a > 0:
			# if a>0, there can be no sum that would be >0. Terminate immediately!
			break

		j = i+1 #the first after a
		k = N-1 #last element in list
		#we're gonna squish it
		while(j<k):
			s=a+inp[j]+inp[k]
			if s == 0:
				results.add(tuple(sorted((origs[i],origs[j],origs[k])))) # found one!

				# handle duplicates
				if inp[j] != inp[k]:
					js = [j]
					ks = [k]
					while j<k and inp[j]==inp[j+1]:
						j+=1
						js.append(j)
					while j<k and inp[k]==inp[k-1]:
						k-=1
						ks.append(k)
					# since the result is a set, I don't need to handle duplicates
					# otherwise I would have had to do (if jj != j and kk != k) or something
					for jj in js:
						for kk in ks:
							if jj != kk:
								results.add(tuple(sorted((origs[i],origs[jj],origs[kk])))) # found one!
				else:
					# j and k elements are equal, just put all the possible combos, the loop will end, since it's squished
					jj = j
					while jj < k:
						kk = k
						while kk > jj:
							results.add(tuple(sorted((origs[i],origs[jj],origs[kk])))) # found one!
							kk-=1
						jj+=1


				k -= 1 #decrement end
			elif s > 0:
				#sum too big. Decrement the end
				k -= 1 #decrement end
			else:
				# s < 0, sum too small, incrememnt the start
				j += 1

	return results

def test_sort(n, L):

	for i in range(n):
		INPUT=tuple(r.randint(-100,100) for _ in range(L))
		print(INPUT)
		srted = quicksort_memory(INPUT)
		print(srted)


INPUT=None#for interactive debugging
def main(array_len=42, repeats=100):
	global INPUT#for interactive debugging

	for _ in range(repeats):
		INPUT = (5,-2,-4,-8,3,-3,1,2,-5,7)

		INPUT=tuple(r.randint(-100,100) for _ in range(array_len))

		bad = bad_algorithm(INPUT)
		good = better_algorithm(INPUT)

		print("INPUT", INPUT)
		# print(quicksort_memory(INPUT))
		print("bad", bad)
		print("good", good)

		print(set(bad)^set(good))
		print(set(bad)-set(good))

		assert bad == good

		print(("#"*100+"\n")*3)


def timing(array_len=100, repeats=10):
	import timeit

	INPUT=tuple(r.randint(-100,100) for _ in range(array_len))

	print("bad algorithm time:", min(timeit.Timer("bad_algorithm(INP)", setup="from __main__ import bad_algorithm; INP={}".format(INPUT)).repeat(repeat=repeats, number=1)))
	print("better algorithm time:", min(timeit.Timer("better_algorithm(INP)", setup="from __main__ import better_algorithm; INP={}".format(INPUT)).repeat(repeat=repeats, number=1)))


if __name__ == '__main__':
	# test_sort(100,100)

	# main(array_len=1000)
	timing(1000,1)