First solution to PA2 with test case added in comments