Better implementation with switch-case