I experienced the same sort of thing a few months ago, I forget exactly what the problem was, but I started writing your general sort of answer with a couple of for/while loops, but I blew them away when I said something like, "Oh you could do this and this in one line", and I started combining lines until I had a one or two line answer.

They offered me the job, but I decided that it wasn't worth moving to Silicon Valley for :-)