When the question shifts, briefly restate the new target before answering. That gives you a mental reset and prevents you from solving the old question after the interviewer has moved on.

Keep the objective at the top of your notes and treat each redirect as a new sub-question within the same overall case.

You do not need to preserve the exact structure of your previous answer if the interviewer has genuinely changed the problem.