Solve problem 4 in Swift
Refactor problem 4 in Swift for better performance