Even more pythonic with some list comp and lambda