Skip to content

Commit a1f9075

Browse files
committed
Added FFT Cooley–Tukey polynomial convolution, segment tree with lazy propagation
1 parent 08d8c6b commit a1f9075

File tree

6 files changed

+451
-0
lines changed

6 files changed

+451
-0
lines changed

Maths/FFT.js

Lines changed: 152 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,152 @@
1+
// Cooley–Tukey FFT (radix-2, iterative) and polynomial/big-integer multiplication
2+
// Exports: fft, ifft, multiplyPolynomials, convolveReal, multiplyBigIntegers
3+
4+
function isPowerOfTwo(n) {
5+
return n && (n & (n - 1)) === 0
6+
}
7+
8+
function nextPowerOfTwo(n) {
9+
let p = 1
10+
while (p < n) p <<= 1
11+
return p
12+
}
13+
14+
function bitReverse(n, bits) {
15+
let rev = 0
16+
for (let i = 0; i < bits; i++) {
17+
rev = (rev << 1) | (n & 1)
18+
n >>= 1
19+
}
20+
return rev
21+
}
22+
23+
function fft(re, im, invert = false) {
24+
const n = re.length
25+
if (!isPowerOfTwo(n)) {
26+
throw new Error('fft input length must be a power of two')
27+
}
28+
if (im.length !== n) throw new Error('re and im lengths must match')
29+
30+
// Bit-reverse permutation
31+
const bits = Math.floor(Math.log2(n))
32+
for (let i = 0; i < n; i++) {
33+
const j = bitReverse(i, bits)
34+
if (i < j) {
35+
;[re[i], re[j]] = [re[j], re[i]]
36+
;[im[i], im[j]] = [im[j], im[i]]
37+
}
38+
}
39+
40+
// Iterative FFT
41+
for (let len = 2; len <= n; len <<= 1) {
42+
const ang = 2 * Math.PI / len * (invert ? 1 : -1)
43+
const wLenRe = Math.cos(ang)
44+
const wLenIm = Math.sin(ang)
45+
46+
for (let i = 0; i < n; i += len) {
47+
let wRe = 1
48+
let wIm = 0
49+
for (let j = 0; j < len / 2; j++) {
50+
const uRe = re[i + j]
51+
const uIm = im[i + j]
52+
const vRe = re[i + j + len / 2]
53+
const vIm = im[i + j + len / 2]
54+
55+
// v * w
56+
const tRe = vRe * wRe - vIm * wIm
57+
const tIm = vRe * wIm + vIm * wRe
58+
59+
// butterfly
60+
re[i + j] = uRe + tRe
61+
im[i + j] = uIm + tIm
62+
re[i + j + len / 2] = uRe - tRe
63+
im[i + j + len / 2] = uIm - tIm
64+
65+
// w *= wLen
66+
const nwRe = wRe * wLenRe - wIm * wLenIm
67+
const nwIm = wRe * wLenIm + wIm * wLenRe
68+
wRe = nwRe
69+
wIm = nwIm
70+
}
71+
}
72+
}
73+
74+
if (invert) {
75+
for (let i = 0; i < n; i++) {
76+
re[i] /= n
77+
im[i] /= n
78+
}
79+
}
80+
}
81+
82+
function ifft(re, im) {
83+
fft(re, im, true)
84+
}
85+
86+
function convolveReal(a, b) {
87+
const need = a.length + b.length - 1
88+
const n = nextPowerOfTwo(need)
89+
90+
const reA = new Array(n).fill(0)
91+
const imA = new Array(n).fill(0)
92+
const reB = new Array(n).fill(0)
93+
const imB = new Array(n).fill(0)
94+
95+
for (let i = 0; i < a.length; i++) reA[i] = a[i]
96+
for (let i = 0; i < b.length; i++) reB[i] = b[i]
97+
98+
fft(reA, imA)
99+
fft(reB, imB)
100+
101+
const re = new Array(n)
102+
const im = new Array(n)
103+
for (let i = 0; i < n; i++) {
104+
// (reA + i imA) * (reB + i imB)
105+
re[i] = reA[i] * reB[i] - imA[i] * imB[i]
106+
im[i] = reA[i] * imB[i] + imA[i] * reB[i]
107+
}
108+
109+
ifft(re, im)
110+
111+
const res = new Array(need)
112+
for (let i = 0; i < need; i++) {
113+
res[i] = Math.round(re[i]) // round to nearest integer to counter FP errors
114+
}
115+
return res
116+
}
117+
118+
function multiplyPolynomials(a, b) {
119+
return convolveReal(a, b)
120+
}
121+
122+
function trimLSD(arr) {
123+
// Remove trailing zeros in LSD-first arrays
124+
let i = arr.length - 1
125+
while (i > 0 && arr[i] === 0) i--
126+
return arr.slice(0, i + 1)
127+
}
128+
129+
function multiplyBigIntegers(A, B, base = 10) {
130+
// A, B are LSD-first arrays of digits in given base
131+
if (!Array.isArray(A) || !Array.isArray(B)) {
132+
throw new Error('Inputs must be digit arrays')
133+
}
134+
const conv = convolveReal(A, B)
135+
136+
// Carry handling
137+
const res = conv.slice()
138+
let carry = 0
139+
for (let i = 0; i < res.length; i++) {
140+
const total = res[i] + carry
141+
res[i] = total % base
142+
carry = Math.floor(total / base)
143+
}
144+
while (carry > 0) {
145+
res.push(carry % base)
146+
carry = Math.floor(carry / base)
147+
}
148+
const trimmed = trimLSD(res)
149+
return trimmed.length === 0 ? [0] : trimmed
150+
}
151+
152+
export { fft, ifft, convolveReal, multiplyPolynomials, multiplyBigIntegers }

Maths/test/FFT.test.js

Lines changed: 59 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,59 @@
1+
import { multiplyPolynomials, multiplyBigIntegers, convolveReal } from '../FFT'
2+
3+
describe('FFT polynomial multiplication', () => {
4+
it('multiplies small polynomials', () => {
5+
const a = [1, 2, 3] // 1 + 2x + 3x^2
6+
const b = [4, 5] // 4 + 5x
7+
expect(multiplyPolynomials(a, b)).toEqual([4, 13, 22, 15])
8+
})
9+
10+
it('convolution matches naive for random arrays', () => {
11+
const a = [0, 1, 0, 2, 3]
12+
const b = [1, 2, 3]
13+
const conv = convolveReal(a, b)
14+
const naive = []
15+
for (let i = 0; i < a.length + b.length - 1; i++) {
16+
let sum = 0
17+
for (let j = 0; j < a.length; j++) {
18+
const k = i - j
19+
if (k >= 0 && k < b.length) sum += a[j] * b[k]
20+
}
21+
naive.push(sum)
22+
}
23+
expect(conv).toEqual(naive)
24+
})
25+
})
26+
27+
describe('FFT big integer multiplication', () => {
28+
function digitsToBigInt(digs, base = 10) {
29+
// LSD-first digits to BigInt
30+
let s = ''
31+
for (let i = digs.length - 1; i >= 0; i--) s += digs[i].toString(base)
32+
return BigInt(s)
33+
}
34+
35+
function bigIntToDigits(n, base = 10) {
36+
if (n === 0n) return [0]
37+
const digs = []
38+
const b = BigInt(base)
39+
let x = n
40+
while (x > 0n) {
41+
digs.push(Number(x % b))
42+
x /= b
43+
}
44+
return digs
45+
}
46+
47+
it('multiplies large integer arrays (base 10)', () => {
48+
const A = Array.from({ length: 50 }, () => Math.floor(Math.random() * 10))
49+
const B = Array.from({ length: 50 }, () => Math.floor(Math.random() * 10))
50+
const prodDigits = multiplyBigIntegers(A, B, 10)
51+
const prodBigInt = digitsToBigInt(A) * digitsToBigInt(B)
52+
expect(prodDigits).toEqual(bigIntToDigits(prodBigInt))
53+
})
54+
55+
it('handles leading zeros and zero cases', () => {
56+
expect(multiplyBigIntegers([0], [0])).toEqual([0])
57+
expect(multiplyBigIntegers([0, 0, 1], [0, 2])).toEqual([0, 0, 0, 2])
58+
})
59+
})

String/SuffixAutomaton.js

Lines changed: 102 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,102 @@
1+
// Suffix Automaton implementation for substring queries
2+
// Provides: buildSuffixAutomaton, countDistinctSubstrings, longestCommonSubstring
3+
4+
class SAMState {
5+
constructor() {
6+
this.next = Object.create(null)
7+
this.link = -1
8+
this.len = 0
9+
}
10+
}
11+
12+
function buildSuffixAutomaton(s) {
13+
const states = [new SAMState()]
14+
let last = 0
15+
16+
for (const ch of s) {
17+
const cur = states.length
18+
states.push(new SAMState())
19+
states[cur].len = states[last].len + 1
20+
21+
let p = last
22+
while (p !== -1 && states[p].next[ch] === undefined) {
23+
states[p].next[ch] = cur
24+
p = states[p].link
25+
}
26+
27+
if (p === -1) {
28+
states[cur].link = 0
29+
} else {
30+
const q = states[p].next[ch]
31+
if (states[p].len + 1 === states[q].len) {
32+
states[cur].link = q
33+
} else {
34+
const clone = states.length
35+
states.push(new SAMState())
36+
states[clone].len = states[p].len + 1
37+
states[clone].next = { ...states[q].next }
38+
states[clone].link = states[q].link
39+
40+
while (p !== -1 && states[p].next[ch] === q) {
41+
states[p].next[ch] = clone
42+
p = states[p].link
43+
}
44+
states[q].link = states[cur].link = clone
45+
}
46+
}
47+
48+
last = cur
49+
}
50+
51+
return { states, last }
52+
}
53+
54+
function countDistinctSubstrings(s) {
55+
const { states } = buildSuffixAutomaton(s)
56+
let count = 0
57+
// State 0 is the initial state; skip it in the sum
58+
for (let v = 1; v < states.length; v++) {
59+
const link = states[v].link
60+
const add = states[v].len - (link === -1 ? 0 : states[link].len)
61+
count += add
62+
}
63+
return count
64+
}
65+
66+
function longestCommonSubstring(a, b) {
67+
// Build SAM of string a, then walk b to find LCS
68+
const { states } = buildSuffixAutomaton(a)
69+
let v = 0
70+
let l = 0
71+
let best = 0
72+
let bestEnd = -1
73+
74+
for (let i = 0; i < b.length; i++) {
75+
const ch = b[i]
76+
if (states[v].next[ch] !== undefined) {
77+
v = states[v].next[ch]
78+
l++
79+
} else {
80+
while (v !== -1 && states[v].next[ch] === undefined) {
81+
v = states[v].link
82+
}
83+
if (v === -1) {
84+
v = 0
85+
l = 0
86+
continue
87+
} else {
88+
l = states[v].len + 1
89+
v = states[v].next[ch]
90+
}
91+
}
92+
if (l > best) {
93+
best = l
94+
bestEnd = i
95+
}
96+
}
97+
98+
if (best === 0) return ''
99+
return b.slice(bestEnd - best + 1, bestEnd + 1)
100+
}
101+
102+
export { buildSuffixAutomaton, countDistinctSubstrings, longestCommonSubstring }
Lines changed: 24 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,24 @@
1+
import { countDistinctSubstrings, longestCommonSubstring } from '../SuffixAutomaton'
2+
3+
describe('Suffix Automaton - distinct substrings', () => {
4+
it('handles empty string', () => {
5+
expect(countDistinctSubstrings('')).toBe(0)
6+
})
7+
8+
it('counts distinct substrings correctly', () => {
9+
expect(countDistinctSubstrings('aaa')).toBe(3)
10+
expect(countDistinctSubstrings('abc')).toBe(6)
11+
expect(countDistinctSubstrings('ababa')).toBe(9)
12+
})
13+
})
14+
15+
describe('Suffix Automaton - longest common substring', () => {
16+
it('finds LCS of two strings', () => {
17+
expect(longestCommonSubstring('xabcdxyz', 'xyzabcd')).toBe('abcd')
18+
expect(longestCommonSubstring('hello', 'yellow')).toBe('ello')
19+
})
20+
21+
it('returns empty when no common substring', () => {
22+
expect(longestCommonSubstring('abc', 'def')).toBe('')
23+
})
24+
})

0 commit comments

Comments
 (0)