Your general approach to solving the problem here is good but there IMO there are some subtle code smells here (see comments). These aren't obvious issues, just subtleties related to code quality and organization that may possibly hurt maintainability in the long run. But don't be discouraged. You'll naturally get better at nuances like naming, organization, and consistency with more experience.