I feel for you. I'd probably be just as distracted as you were if this happened in the middle of an interview, but in the off case I weren't, I'd probably talk the interviewers' ear off. What they seem to not understand is, this kind of correction is natural and perfectly normal when you're writing code. I do it all the time - I think of a small piece of code to write, I write it, and before I'm done I notice obvious improvements, which I do on the fly.

It's not "premature optimization", as it costs me nothing (on the contrary, not doing it will cost me my focus / peace of mind), improves performance of the code - or rather, avoids stupidity - and often improves readability, as the code is now closer to a "pure" description of what it's supposed to do.

I've long held that a big part of why software today sucks so badly is because people don't do this, or likely can't do this. They write stupid code, and never bother to learn how to write non-stupid version of the same code by default.