Found a way to simplify the traverse function