Besides the efficient solution by alexwlchan, I'd like to also make some suggestions.
Creating the list
You use
lst = []
for i in range(1, n + 1):
lst.append((i** 2 + i)//2)
Shorter and more efficient would be
lst = [(i ** 2 + i) // 2 for i in range(1, n+1)]
(But you actually don't need to store the list.)
Checking positions
Notice how you check lst.index
. That makes that the statement in the inner loop is only true for at most 2 entries in the entire list. But lst.index
needs to search through the list, which is not that fast.
for i in lst:
for j in lst:
if i*i + j*j == n and ((lst.index(i) == lst.index(j)+1)
or (lst.index(i) == lst.index(j)-1)):
return True
else:
continue
Either i < j or i > j. If i > j, swapping the roles of i and j would give i < j, and (thus) would have shown up earlier in the iteration. So you only need to check the first condition, as that's always the one to be triggered.
for i in lst:
for j in lst:
if i*i + j*j == n and lst.index(i) == lst.index(j)+1:
return True
else:
continue
Now this only helps a little bit. What we actually want is to get rid of the inner loop.
for idx, i in enumerate(lst):
try:
j = lst[idx + 1]
except IndexError:
continue
if i*i + j*j == n:
return True
Now, I'm not a fan of the try/except
here. But that's because we used a list.
Getting rid of the list
def t(n):
return (n ** 2 + n) // 2
def is_consecutive_triangular_square_sum(n):
# Blatantly stealing name from alexwlchan
for i in range(1, n + 1):
if t(i)**2 + t(i+1)**2 == n:
return True
return False
This has the downside of calling t
multiple times for each i
. But we could get rid of that if performance is paramount (which it probably isn't, but let's presume it is!)
Optimised
def is_consecutive_triangular_square_sum(n):
ts = ((i**2 + i)//2 for i in range(1, n+1))
t_prev = next(ts)
for t_next in ts:
if t_prev * t_prev + t_next * t_next == n:
return True
t_prev = t_next
return False
And short-circuiting to stop early.
def is_consecutive_triangular_square_sum(n):
ts = ((i**2 + i)//2 for i in range(1, n+1))
t_prev = next(ts)
for t_next in ts:
squared_sum = t_prev * t_prev + t_next * t_next
if squared_sum == n:
return True
elif squared_sum > n:
# At this point, all squares will be larger anyway
# since the squared-sum is ever-increasing.
break
t_prev = t_next
return False