Again the python std lib saved me some lines of code. itertools to the rescue and basically we are done.
itertools