【PTA数据结构 | C语言版】MillerRabin 素数测试算法

本专栏持续输出数据结构题目集，欢迎订阅。

题目

本题请你实现 Miller-Rabin 素数测试算法，其原理在“101 计划”核心教材《数据结构》第 16 章中有介绍。

输入格式：
输入在一行中给出两个正整数：n（≤1000）为待判定的数字、k（≤10）为测试次数。

输出格式：
若 n 为素数，则在一行中输出 1，否则输出 0。

输入样例 1：
561 3
输出样例 1：
0
输入样例 2：
17 3
输出样例 2：
1

代码

#include <stdio.h>
#include <stdlib.h>
#include <time.h>

// 计算 (a * b) % mod，防止溢出
long long mul_mod(long long a, long long b, long long mod) {
    long long result = 0;
    a = a % mod;
    while (b > 0) {
        if (b % 2 == 1) {
            result = (result + a) % mod;
        }
        a = (a * 2) % mod;
        b /= 2;
    }
    return result % mod;
}

// 计算 (base^exponent) % mod，使用快速幂算法
long long pow_mod(long long base, long long exponent, long long mod) {
    long long result = 1;
    base = base % mod;
    while (exponent > 0) {
        if (exponent % 2 == 1) {
            result = mul_mod(result, base, mod);
        }
        exponent = exponent >> 1;  // 等价于exponent /= 2
        base = mul_mod(base, base, mod);
    }
    return result;
}

// Miller-Rabin单次测试
int miller_rabin_test(long long n, long long a) {
    // 处理n的偶数情况
    if (n % 2 == 0) {
        return n == 2;
    }
    
    // 将n-1表示为d*2^s
    long long d = n - 1;
    int s = 0;
    while (d % 2 == 0) {
        d /= 2;
        s++;
    }
    
    // 计算a^d mod n
    long long x = pow_mod(a, d, n);
    
    // 如果x == 1或x == n-1，可能是素数
    if (x == 1 || x == n - 1) {
        return 1;
    }
    
    // 进行s-1次平方操作
    for (int j = 1; j < s; j++) {
        x = pow_mod(x, 2, n);
        if (x == n - 1) {
            return 1;
        }
    }
    
    // 经过所有测试都不满足，一定不是素数
    return 0;
}

// Miller-Rabin素数测试主函数
int is_prime(long long n, int k) {
    // 处理小于2的情况
    if (n <= 1) {
        return 0;
    }
    // 处理2和3的情况
    if (n <= 3) {
        return 1;
    }
    
    // 初始化随机数生成器
    srand(time(0));
    
    // 进行k次测试
    for (int i = 0; i < k; i++) {
        // 生成1到n-1之间的随机数a
        long long a = rand() % (n - 1) + 1;
        
        // 如果有一次测试失败，则n不是素数
        if (!miller_rabin_test(n, a)) {
            return 0;
        }
    }
    
    // 所有测试通过，n很可能是素数
    return 1;
}

int main() {
    int n, k;
    scanf("%d %d", &n, &k);
    
    printf("%d\n", is_prime(n, k));
    
    return 0;
}