For the record, the full comment:
For the record, usually we're against updating the code in the question after answers have been received. It creates a mess if more answers come in, reviews of all different kinds of versions make it hard to follow what's going on. Please see what you may and may not do after receiving answers. However, since the author of the only answer so far is ok with it and the changes are minor, it's fine this time. Just, try not to do it again, ok? Thanks :-)
There's limited community consensus about when to and when not to rollback. The clear cases are easy, but there's a lot of not-so-clear cases.
In general, it's simply not accepted that code is modified after answers come in. Why? So every reviewer gets the same version of the code to review. Which is important, so if anyone finds your question after a year all reviews are about the same version of the code. It tends to get a mess otherwise.
There are however a couple of examples where allowing modification of a piece of code is "the lesser evil". Or where it's simply not worth the hassle to roll it back. I've been involved with a lot of rollbacks, and too many of them end in complaints, "recidivism" or other situations which require extra effort to resolve. Or when the question was actually off-topic, but OP and answerer teamed together to fix both the question and the answer. At which point a rollback would be pointless. So often enough cases are condoned, tolerated. That doesn't make them right.
The overall advice is not to make edits after answers come in. Adding a button for this very specific situation would just about legalise it under certain conditions and send the wrong message altogether. Would go against site policy.
You don't just post a question. You post a piece of code that does something, someone tells you it can be improved (and how) and you learned something. Important here is that other people learn something when reading both question and answer too, otherwise we might just as well destroy every question older than X weeks. We have to keep readability for those people in mind too, and sometimes that means keeping the mistakes intact and not allowing further edits when answers have arrived.
So, please no.