Refactor using each_with_index