Prime Arrangements

easy math combinatorics sieve

Problem

Return the number of permutations of the numbers 1 to n so that the prime numbers are at prime indices (1-indexed). Because the answer may be large, return it modulo 10^9 + 7.

Inputn = 5

Output12

Primes among 1..5 are {2, 3, 5} → 3 of them, and prime indices are {2, 3, 5} → 3 slots. The 3 primes can fill the 3 prime slots in 3! = 6 ways and the 2 non-primes fill the 2 non-prime slots in 2! = 2 ways, so 6 × 2 = 12.

def num_prime_arrangements(n):
    MOD = 10**9 + 7

    def is_prime(x):
        if x < 2:
            return False
        d = 2
        while d * d <= x:
            if x % d == 0:
                return False
            d += 1
        return True

    p = sum(1 for x in range(1, n + 1) if is_prime(x))
    ans = 1
    for k in range(2, p + 1):
        ans = ans * k % MOD
    for k in range(2, n - p + 1):
        ans = ans * k % MOD
    return ans

function numPrimeArrangements(n) {
  const MOD = 1000000007n;
  const isPrime = (x) => {
    if (x < 2) return false;
    for (let d = 2; d * d <= x; d++) if (x % d === 0) return false;
    return true;
  };
  let p = 0;
  for (let x = 1; x <= n; x++) if (isPrime(x)) p++;
  let ans = 1n;
  for (let k = 2; k <= p; k++) ans = (ans * BigInt(k)) % MOD;
  for (let k = 2; k <= n - p; k++) ans = (ans * BigInt(k)) % MOD;
  return Number(ans);
}

class Solution {
    public int numPrimeArrangements(int n) {
        long MOD = 1000000007L;
        int p = 0;
        for (int x = 1; x <= n; x++) if (isPrime(x)) p++;
        long ans = 1;
        for (int k = 2; k <= p; k++) ans = ans * k % MOD;
        for (int k = 2; k <= n - p; k++) ans = ans * k % MOD;
        return (int) ans;
    }

    private boolean isPrime(int x) {
        if (x < 2) return false;
        for (int d = 2; (long) d * d <= x; d++) if (x % d == 0) return false;
        return true;
    }
}

class Solution {
public:
    int numPrimeArrangements(int n) {
        const long long MOD = 1000000007LL;
        auto isPrime = [](int x) {
            if (x < 2) return false;
            for (int d = 2; (long long) d * d <= x; d++)
                if (x % d == 0) return false;
            return true;
        };
        int p = 0;
        for (int x = 1; x <= n; x++) if (isPrime(x)) p++;
        long long ans = 1;
        for (int k = 2; k <= p; k++) ans = ans * k % MOD;
        for (int k = 2; k <= n - p; k++) ans = ans * k % MOD;
        return (int) ans;
    }
};

Explanation

The key insight is that the primes can only sit in prime-numbered slots, and the non-primes can only sit in the leftover slots. These two groups never mix, so we can count each group's arrangements separately and multiply.

First we count how many numbers in 1..n are prime — call it p. Conveniently, the number of prime positions (indices 2, 3, 5, 7, ...) is exactly the same p, so the primes fill those slots in p! ways.

The remaining n - p non-prime numbers fill the n - p non-prime slots in (n - p)! ways. The total is p! × (n - p)!. Because that product can get huge, we take it modulo 10^9 + 7 as we multiply.

The is_prime helper just tries dividing by every d up to √x; if nothing divides evenly, x is prime.

Example: n = 5. Primes are {2, 3, 5} so p = 3. Answer = 3! × 2! = 6 × 2 = 12.

Time: O(n · √n) Space: O(1)