Reduced problem1 to two lines & made it more pythonic