A function to compute the first N primes

So I wanted to burn a few minutes before actual deep work starts, went to a random leetcode problem, group anagrams. There are (obviously) many ways to solve this. But we can make it fun. From the fundamental theorem of arithmetic (I think that’s what it’s called), every integer greater than 1 has a unique prime factorization. So the product of primes is always unique. We can use this: map each letter to a prime ($a=2, b=3, c=5, \ldots$), multiply them together for each word, and words with the same product are anagrams. “eat” and “tea” both give $2 \times 5 \times 71 = 710$. Same product, same group.

One caveat with the prime product approach is that it can overflow. The larger primes in the mapping go up to 101, and multiplying several of those together for longer strings will blow past long long limits. You might think to mod the product by a large prime (the classic $10^9 + 7$), but that destroys the uniqueness guarantee, which is the whole reason we used prime products in the first place. Two different products can collide to the same value mod $M$, grouping non-anagrams together. LeetCode would probably accept it since its test cases are unlikely to hit a collision, but the solution is technically wrong. The simpler and more common approach is to just sort each word’s characters and use that sorted string as the key.

I am going to omit the full solution for brevity, but what I find more fun is a thought I had when thinking about the problem: “if I want a function that takes an integer $N$ and returns the first $N$ primes, how should I write it?”. There are many ways, especially if you look at competitive programmers’ solutions with all their optimizations. For the sake of this bit, let’s keep it simple.

std::vector<int> first_n_primes(int n) {
    std::vector<int> primes;
    primes.reserve(n);
    for (int x = 2; primes.size() < (size_t)n; ++x) {
        bool ok = true;
        for (int p : primes) {
            if (p * p > x) break;
            if (x % p == 0) { ok = false; break; }
        }
        if (ok) primes.push_back(x);
    }
    return primes;
}

The code is short, but why and how does it works?

First, let’s see why we only check agains primes. We know that a prime is a number greater than 1 whose only divisors are 1 and itself. A composite (non-prime) is, by definition, a product of smaller primes. So if some composite $c$ divides $x$, then the prime factors of $c$ also divide $x$, and those prime factors are smaller than $c$. Checking against primes is sufficient because any composite divisor would just be redundant, its prime factors already catch it.

In other words: if $x$ is divisible by 6, it is divisible by 2 and 3. We will find 2 (or 3) before we ever reach 6.

Now, why stop at $p^2 > x$? This is an optimization; in particular if $x$ has a factor larger than $\sqrt{x}$, it must also have a corresponding factor smaller than $\sqrt{x}$. This follows directly from multiplication: if $x = a \cdot b$ and both $a > \sqrt{x}$ and $b > \sqrt{x}$, then $a \cdot b > x$, a contradiction. So at least one factor must be $\leq \sqrt{x}$.

The check p * p > x is equivalent to $p > \sqrt{x}$ but avoids floating point. If no prime up to $\sqrt{x}$ divides $x$, nothing will, and $x$ is prime.

Walk through $x = 37$: primes collected so far are $[2, 3, 5, 7, \ldots]$. We check $2$ ($4 < 37$, $37 \bmod 2 \neq 0$), $3$ ($9 < 37$, $37 \bmod 3 \neq 0$), $5$ ($25 < 37$, $37 \bmod 5 \neq 0$), $7$ ($49 > 37$, stop). No divisor found, so 37 is prime.

And why does the outer loop start at 2? Well we all know $2$ is the smallest prime. On the first iteration, primes is empty, so the inner loop does not execute, ok stays true, and $2$ gets added. Basically this bootstraps the whole process in the sense that once we have $2$, we can check $3$ against it, then $4$, and so on. Each new prime we discover extends our set of trial divisors.

Of course, for the actual group anagrams solution, if you (somehow) memorized the first 26 prime numbers you can just hardcode them and skip the computation entirely (although it’s not that much for small $N$ anyways):

constexpr int primes[26] = {
    2, 3, 5, 7, 11, 13, 17, 19, 23, 29, 31, 37, 41,
    43, 47, 53, 59, 61, 67, 71, 73, 79, 83, 89, 97, 101
};