How to pick initial values for Shor's algorithm?

Question

From my understanding Shor's algorithm can be summarised as follows...

Task: given a positive integer $N$ find another positive integer $s$ such that $s$ divides $N$.

Steps (informal)

1. Choose an integer $a$ uniformly at random from the interval $[1, N]$
2. Compute $b = \gcd(a, N)$. If $b > 1$, then $b$ divides $N$. Otherwise continue to 3.
3. Determine the period of the function $f(x) = a^x \,\text{mod} \,N$. In other words, find the least positive integer $r$ such that $f(x+r)=f(x)$. Note that this is the part that uses the quantum computer. This is typically done with quantum phase estimation.
4. Now if $a^{r/2}$ is an integer evaluate $s=\gcd(a^{r/2}-1, N)$. If $a^{r/2}$ is not an integer our algorithm fails. Likewise if $s=1$, our algorithm hasn't worked and we try need to try again. If $s \neq 1$, then $s$ is a factor of $N$ and our algorithm has worked.

Note that the steps above give you a single non-trivial factor of $N$. In the textbook version of Shor's algorithm, we usually have to repeat the above to get all the factors of $N$.

My question is mostly about step 1 where we choose $a$.

A key part of the algorithm is this difference of two squares factorisation.

$$ a^r -1 = (a^{r/2} -1)(a^{r/2} + 1) = 0 \,\, \text{mod} \, N $$

Now if $a^{r/2}$ is not an integer, then we're in trouble as $\gcd(a^{r/2}-1, N)$ is not defined. This is usually the case when our quantum computer gives us an odd number for the period $r$. However in the edge case where $r$ is odd but $a$ is a perfect square the algorithm can still succeed. This is becuase if $a=j^2$ for some positive integer $j$, then $a^{r/2}=(j^2)^{r/2} =j^r$.

I'm asking the following... What is the reason for choosing $a$ uniformly at random from the positive integers less than $N$? Have people considered other distributions from which to sample $a$?

A trivial way to choose $a$ would be to sample from the set of all of the perfect squares less than $N$. This would mean that $a^{r/2}$ is always an integer regardless of the value of the period $r$. Surely it cannot be this simple and I must have misunderstood something or missed some other requirements for the algorithm to succeed. Perhaps sampling from contrived distributions makes it less likely to find non-trivial factors of $N$? I'm not sure.

Example

Consider the basic example of factoring $N=21$. Suppose our random choice in step 1. gives us $a=4$.

Now we have to determine the period of the function $f(x) = 4^x \, \text{mod} \, 21$. We can determine that the period is $r=3$ by evaluating $f(x)$ on a range of inputs.

Now normally the fact that $r$ is odd would cause problems. However in this case $a^{r/2}=4^{3/2}=8$ (an integer as $4$ is a perfect square).

Now we see that $\gcd(a^{r/2}-1,\, N)=\gcd(7,\, 21)=7$. Clearly $7$ divides $21$. So we have found that $21 = 7 \times 3$.

Answer 1

Shor proves in [Shor94] [Shor97] that if $N$ is a positive odd integer that has $k \ge 2$ distinct prime factors, and if $a$ is selected uniformly at random from $\mathbb Z_N^*$ as in your steps 1–2, then the probability that the order $r$ of $a$ is even and that $a^{r/2} \not\equiv -1 \:\: (\text{mod } N)$ is at least $1 - 2^{1-k} \ge 1/2$.

The uniform selection of $a$ from $\mathbb Z_N^*$ greatly facilitates the derivation of said proof.

In practice, it is possible to select $a$ in other ways, but depending on how $a$ is selected it may then be considerably more difficult to prove that the algorithm works for general $N$. In the literature, there are examples of people considering other ways of selecting $a$ for special-form $N$. One such example that comes to mind is [GLMS15]. Another example, that is somewhat related to the discussion in the comments on your question, is [Leander02].

The above having been said, note also that by replacing your step 4 with a slightly more advanced classical post-processing step, it is possible to guarantee finding the complete factorization of $N$ with probability very close to one, see my work [E21b] (and also my work [E24] for the complete algorithm in your steps 1–4). In short, with a uniform selection of $a$ from $\mathbb Z_N^*$, we expect to have to run the quantum order-finding part only once to completely factor $N$ into all of its prime factors. There is no need for special treatment of odd $r$, etc.