How does the uncomputation step work in the Grover-Rudolph scheme to prepare $\sum_i\sqrt{p_i}|i\rangle$?

Question

In https://arxiv.org/abs/quant-ph/0208112, the authors discuss a scheme to, given a discrete probability distribution $\mathbf p\equiv (p_i)_i$, under some assumptions on $\mathbf p$, prepare the superposition state $$|\psi(\mathbf p)\rangle = \sum_i \sqrt{p_i}|i\rangle.$$ To have a toy example in mind and concretise the discussion, say we want to prepare the two-qubit state $$\sum_{i=0}^3\sqrt{p_i}|i\rangle\equiv \sqrt{p_0}|00\rangle+\sqrt{p_1}|01\rangle +\sqrt{p_2}|10\rangle+ \sqrt{p_3}|11\rangle.$$ As also discussed in How does the induction step in the Grover-Rudolph scheme to prepare superpositions from probabilities work?, the gist of the scheme is to iteratively "refine" the superposition. In other words, you start from $|0\rangle$, evolve it to $\sqrt{p_0+p_1}|0\rangle+\sqrt{p_2+p_3}|1\rangle$, and then implement another evolution such that $$\sqrt{p_0+p_1}|0\rangle \to \sqrt{p_0}|00\rangle + \sqrt{p_1}|01\rangle, \\ \sqrt{p_2+p_3}|1\rangle \to \sqrt{p_2}|10\rangle + \sqrt{p_3}|11\rangle.$$ To implement these evolutions, the authors propose a two-step procedure:

1. Evolve $|0\rangle$ to $|0\rangle|\theta_0\rangle$ for some suitably defined $\theta_0$. This you can do as long as computing $\theta_0$ is efficient classically.
2. Use $\theta_0$ as a control to perform a rotation on another ancillary qubit. Suitably choosing $\theta_0$, you can thus way implement the evolution $$|0\rangle|\theta_0\rangle\to |0\rangle|\theta_0\rangle\frac{\sqrt{p_0}|0\rangle+\sqrt{p_1}|1\rangle}{\sqrt{p_0+p_1}}.$$

The same procedure is applied separately to $|1\rangle$, resulting in $$|1\rangle|\theta_1\rangle\to |1\rangle|\theta_1\rangle\frac{\sqrt{p_2}|0\rangle+\sqrt{p_3}|1\rangle}{\sqrt{p_2+p_3}}.$$ Overall, this process gives us the evolution $$\sqrt{p_0+p_1}|0\rangle+\sqrt{p_2+p_3}|1\rangle \to |0\rangle|\theta_0\rangle(\sqrt{p_0}|0\rangle+\sqrt{p_1}|1\rangle) + |1\rangle|\theta_1\rangle(\sqrt{p_2}|0\rangle+\sqrt{p_3}|1\rangle).$$ That's almost the target superposition, except for the lingering ancillary qubits $|\theta_i\rangle$, which are however entangled with the rest of the qubits and therefore cannot just be ignored.

The authors just write that "we uncompute the register containing $|\theta_i\rangle$ to leave us with the desired state. How does this computation step work precisely, and what's its cost?

Answer 1

The uncomputation is performed by simply running the same unitary used to generate the $|\theta_i\rangle$ states in reverse.

We could describe what's going on more abstractly as follows:

1. Start with some $\alpha|0\rangle+\beta|1\rangle$
2. Implement a unitary $U$ such that $U|i\rangle|0\rangle=|i\rangle|\theta_i\rangle$, $i=0,1$. Here $\theta_i\in\mathbb{Z}_2^n$ if the ancilla register used to store the rotation angles $\theta_i$ uses $n$ qubits.
3. Run the controlled rotations. If the corresponding unitary is $R$, then by definition this is an operation such that $$R|i\rangle|\theta_i\rangle|0\rangle = |i\rangle|\theta_i\rangle |\psi_i\rangle$$ for some states $|\psi_i\rangle\equiv R_i|0\rangle$. The overall state we now have is $$RU(\alpha|0\rangle+\beta|1\rangle) = \alpha|0\rangle|\theta_0\rangle|\psi_0\rangle + \beta|1\rangle|\theta_1\rangle|\psi_1\rangle.$$
4. For the uncomputation, we just use $U^\dagger$. From the definition of $U$, it immediately follows that $U^\dagger|i\rangle|\theta_i\rangle=|i\rangle|0\rangle$, and thus $$U^\dagger R U (\alpha|0\rangle+\beta|1\rangle)=\alpha|0\rangle|0\rangle|\psi_0\rangle+\beta|1\rangle|0\rangle|\psi_1\rangle \simeq \alpha |0\rangle|\psi_0\rangle+\beta|1\rangle|\psi_1\rangle.$$