Implemented solution for w2 in Python