魔法师の高塔

快速数论变换

在快速傅里叶变换中,由于复数原因,在计算大数乘法时精度有影响,而且复数计算比实数代价要大。所以寻找一种类似单位根性质的结构来替换单位根,这就是快速数论变换。

具体原理参见数论相关内容。快速数论变换具体结构与快速傅里叶变换类似。Rust实现如下。

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
use std::ops::{Add, Mul};
#[derive(Debug)]
struct Vector {
value: Vec<i64>,
}
impl Vector {
fn new(value: Vec<i64>) -> Vector {
Vector { value: value }
}
}
impl Add for Vector {
type Output = Vector;
fn add(self, vec: Vector) -> Vector {
let mut result = Vec::new();
for i in 0..self.value.len() {
result.push(self.value[i] + vec.value[i]);
}
Vector::new(result)
}
}
impl Mul for Vector {
type Output = Vector;
fn mul(self, vec: Vector) -> Vector {
let mut result = Vec::new();
for i in 0..self.value.len() {
result.push(self.value[i] * vec.value[i]);
}
Vector::new(result)
}
}
static P: i64 = (479 << 21) + 1;
static N: i64 = 1 << 18;
static G: i64 = 3;
fn modexp_rec(a: i64, b: i64, n: i64) -> i64 {
// 快速幂取模
if b == 0 {
1i64
} else if b == 1 {
a % n
} else {
let mut t = 1;
t = modexp_rec(a, b >> 1, n);
t = t * t % n;
if b & 0x1 == 1 {
t = t * a % n;
}
t
}
}
fn ntt(x: Vector, res: bool) -> Vector {
if x.value.len() == 1 {
x
} else {
let mut a0 = Vector::new(Vec::new());
let mut a1 = Vector::new(Vec::new());
let mut result = Vector::new(Vec::new());
for _ in 0..x.value.len() {
result.value.push(0);
}
for i in 0..x.value.len() {
if i % 2 == 0 {
a0.value.push(x.value[i]);
} else {
a1.value.push(x.value[i]);
}
}
a0 = ntt(a0, res);
a1 = ntt(a1, res);
for i in 0..x.value.len() / 2 {
let mut wn = modexp_rec(G, (P - 1) * (i as i64) / (x.value.len() as i64), P);
if res == true {
wn = modexp_rec(wn, P - 2, P);
}
result.value[i] += (a0.value[i] + a1.value[i] * wn) % P;
result.value[i + x.value.len() / 2] += (a0.value[i] - a1.value[i] * wn) % P + P;
}
for i in 0..x.value.len() {
result.value[i] = result.value[i] % P;
}
result
}
}
fn ntt_1(x: Vector) -> Vector {
let mut result = ntt(x, true);
for i in 0..result.value.len() {
result.value[i] = result.value[i] / (result.value.len() as i64);
}
result
}
fn change(x: Vec<i64>) -> Vec<i64> {
let mut result1 = Vec::new();
for i in x {
result1.push(i);
}
for i in 0..result1.len() {
if result1[i] > 9 {
println!("{:?}", result1[i]);
result1[i + 1] += result1[i] / 10;
result1[i] = result1[i] % 10;
}
}
result1
}
fn main() {
let a0 = Vector::new(vec![3, 4, 2, 2, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0]);
let a1 = Vector::new(vec![1, 0, 3, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0]);
let a0result = ntt(a0, false);
let a1result = ntt(a1, false);
// let result = a0result * a1result;
println!("{:?}", 1301 * 2243);
let ntt = ntt_1(a0result * a1result);
println!("{:?}", change(ntt.value));
// println!("{:?}", nft2vec(ntt.value));
}

这个实现非常dirty,因为我想实现向量的运算,然而Vec这个类型和Add这些特质都是独立的,所以无法在这个文件中实现。我应该找找有没有向量计算库。