I think your optimization works and is even a security improvement over the original, because it avoids the attack I have described here. I think it is quite similar to what I proposed at the end of that post. I proposed to follow the path in the forward direction and you are doing it in the reverse direction.
You can find some extra optimizations in my followup post here that you might be able to use.