The solution above is incorrect - I have given a counter example elsewhere… I think that perhaps this is a good place to mention something that I think could be very helpful for everyone who takes the exam tomorrow:
The burden of proof is on you.
It is not enough to find a solution that you think is probably correct - one of the points of the course is being able to find a simple solution, and prove that it is correct. In fact, we prefer a sub-optimal solution which you can prove is correct than a solution which is faster and "probably works".
In this case, once you modify the function pi, you are left with quite a difficult job to prove that everything works… you've modified the algorithm - the proofs from class no longer apply, so essentially you have to re-prove pretty much everything, or give a justification as to why the proofs still hold.
In general, try to modify existing algorithms as little as possible, and use them as black boxes whenever possible… in this sense, at least, your solution is not "alot simpler", but in fact, much more complicated!