Refactor solution for problem C2 some more