Triangle Square Numbers generator from Standupmaths' puzzle

triangle_square_number.py

import math
import itertools

from collections import OrderedDict

"""
Output:
(36.0, 6, 8)
(1225.0, 35, 49)
(41616.0, 204, 288)
(1413721.0, 1189, 1681)
(48024900.0, 6930, 9800)
(1631432881.0, 40391, 57121)
(55420693056.0, 235416, 332928)
(1882672131025.0, 1372105, 1940449)
(63955431761796.0, 7997214, 11309768)
(2172602007770041.0, 46611179, 65918161)
(3.1843510970040004e+16, 178447502, 252362877)
(7.38045128324196e+16, 271669860, 384199200)
(2.0986896474365978e+17, 458114576, 647871846)
(3.950937882553385e+17, 628564864, 888924955)
(6.384252131255831e+17, 799015152, 1129978064)
(9.398632393543936e+17, 969465440, 1371031173)
(1.29940786694177e+18, 1139915728, 1612084282)
(1.4493594148750336e+18, 1203893440, 1702562430)
(1.7170590958877123e+18, 1310366016, 1853137391)
(1.888820682692938e+18, 1374343728, 1943615539)
(2.386388551869408e+18, 1544794016, 2184668648)
(2.7266813578552945e+18, 1651266592, 2335243609)
(3.165630982219104e+18, 1779222016, 2516199905)
(3.5558440942980465e+18, 1885694592, 2666774866)
(4.4949198051087985e+18, 2120122592, 2998306123)
(4.957726042160948e+18, 2226595168, 3148881084)
(5.44320509809325e+18, 2333067744, 3299456045)
(5.54390849028755e+18, 2354550592, 3329837380)
(6.056635033432756e+18, 2461023168, 3480412341)
(6.592034395458114e+18, 2567495744, 3630987302)
(7.265456999072564e+18, 2695451168, 3811943598)
(7.850776667190978e+18, 2801923744, 3962518559)
(8.458769154189542e+18, 2908396320, 4113093520)
(9.219431913291842e+18, 3036351744, 4294049816)
(9.877344706383462e+18, 3142824320, 4444624777)
(1.0557930318355235e+19, 3249296896, 4595199738)
(1.3044234733875302e+19, 3611680320, 5107687291)
(1.3824660958033187e+19, 3718152896, 5258262252)
(1.4627760001071223e+19, 3824625472, 5408837213)
(1.7531043495183139e+19, 4187008896, 5921324766)
(1.8433983150407287e+19, 4293481472, 6071899727)
(1.9359595624511586e+19, 4399954048, 6222474688)
(2.0307880917496037e+19, 4506426624, 6373049649)
(2.127883902936064e+19, 4612899200, 6523624610)
(2.167707792980509e+19, 4655864896, 6584387280)
(2.267985819721535e+19, 4762337472, 6734962241)
(2.3705311283505762e+19, 4868810048, 6885537202)
"""


class LimitedSizeDict(OrderedDict):
  def __init__(self, *args, **kwds):
    self.size_limit = kwds.pop("size_limit", None)
    OrderedDict.__init__(self, *args, **kwds)
    self._check_size_limit()

  def __setitem__(self, key, value):
    OrderedDict.__setitem__(self, key, value)
    self._check_size_limit()

  def _check_size_limit(self):
    if self.size_limit is not None:
      while len(self) > self.size_limit:
        self.popitem(last=False)

# n^^2 = (m(m+1))/2

def triangle_gen():
    for n in itertools.count(start=2):
        yield n, (n * (n + 1)) / 2

def square_gen():
    for n in itertools.count(start=2):
        yield n, n**2


square_dict = LimitedSizeDict(size_limit=3)
square_dict_max = 0
squares = square_gen()
def get_square_root(n):
    global square_dict
    global square_dict_max
    while n > square_dict_max:
        s, square = next(squares)
        square_dict[square] = s
        square_dict_max = square
    return square_dict.get(n)


def square_triangle_gen():
    for t_root, triangle in triangle_gen():
        s_root = get_square_root(triangle)
        if s_root:
            yield triangle, s_root, t_root



def main():
    stgen = square_triangle_gen()
    for i in range(100):
        print(next(stgen))


if __name__ == "__main__":
    main()