neetcode/medium/3sum.py

71 lines
2.1 KiB
Python

#!/usr/bin/python
# yigid balaban <fyb@fybx.dev>
# neetcode 2024
# 3sum
def solution0(nums: list[int]) -> list[list[int]]:
if len(nums) < 3:
return []
if len(nums) == 3 and sum(nums) == 0:
return [nums]
nums = sorted(nums)
triples = []
def comb(what: list[int], head: int):
mtriples = []
for p0 in range(0, len(what) + 1):
for p1 in range(p0, len(what)):
if p1 != p0 and head + what[p0] + what[p1] == 0:
triple = sorted([head, what[p0], what[p1]])
if triple not in triples and triple not in mtriples:
mtriples.append(triple)
return mtriples
for p in range(len(nums) - 3):
triples.extend(comb(nums[p + 1:], nums[p]))
return triples
def solution(nums: list[int]) -> list[list[int]]:
if len(nums) < 3:
return []
if len(nums) == 3:
if sum(nums) == 0: return [nums]
else: return []
nums.sort()
triplets: list[list[int]] = []
for p in range(len(nums)):
if p == len(nums) - 3:
if nums[-1] + nums[-2] + nums[-3] == 0:
triplets.append([nums[-1], nums[-2], nums[-3]])
break
pl = p + 1
pr = len(nums) - 1
while pl != pr:
print(p, pl, pr)
s = nums[p] + nums[pl] + nums[pr]
if s < 0:
pl += 1
elif s > 0:
pr -= 1
else:
triplets.append([nums[p], nums[pl], nums[pr]])
break
p = 1
while p < len(triplets):
if sorted(triplets[p]) == sorted(triplets[p - 1]):
triplets.pop(p)
else: p += 1
return triplets
print(solution([1,-1,-1,0]), solution([1,-1,-1,0]) == [[-1, 0, 1]])
print(solution([0,0]) == [])
print(solution([-1,0,1,2,-1,-4]), solution([-1,0,1,2,-1,-4]) == [[-1,-1,2],[-1,0,1]])
print(solution([0,1,1]) == [])
print(solution([0,0,0]) == [[0,0,0]])
print(solution([0,0,0,0]) == [[0,0,0]], solution([0,0,0,0]))
print(solution([0,0,0,0,0]))