Broke problem into 2 python programs