Last active
March 22, 2024 17:48
-
-
Save 8mist/affe585ccb7b2969fed83ff0e90d7c14 to your computer and use it in GitHub Desktop.
[1kyu] Codewars - Express number as sum of four squares
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| const ROBUST_TEST_SET = [ | |
| [2047n, [2n]], | |
| [1373653n, [2n, 3n]], | |
| [9080191n, [31n, 73n]], | |
| [25326001n, [2n, 3n, 5n]], | |
| [3215031751n, [2n, 3n, 5n, 7n]], | |
| [4759123141n, [2n, 7n, 61n]], | |
| [1122004669633n, [2n, 13n, 23n, 1662803n]], | |
| [2152302898747n, [2n, 3n, 5n, 7n, 11n]], | |
| [3474749660383n, [2n, 3n, 5n, 7n, 11n, 13n]], | |
| [341550071728321n, [2n, 3n, 5n, 7n, 11n, 13n, 17n]], | |
| [3825123056546413051n, [2n, 3n, 5n, 7n, 11n, 13n, 17n, 19n, 23n]], | |
| [ | |
| 18446744073709551616n, | |
| [2n, 3n, 5n, 7n, 11n, 13n, 17n, 19n, 23n, 29n, 31n, 37n], | |
| ], | |
| ]; | |
| const RABIN_EXCEPTIONS = { | |
| 5: [2n, 1n, 0n, 0n], | |
| 10: [3n, 1n, 0n, 0n], | |
| 13: [3n, 2n, 0n, 0n], | |
| 34: [5n, 3n, 0n, 0n], | |
| 37: [6n, 1n, 0n, 0n], | |
| 58: [7n, 3n, 0n, 0n], | |
| 61: [6n, 5n, 0n, 0n], | |
| 85: [9n, 2n, 0n, 0n], | |
| 130: [11n, 3n, 0n, 0n], | |
| 214: [14n, 3n, 3n, 0n], | |
| 226: [15n, 1n, 0n, 0n], | |
| 370: [19n, 3n, 0n, 0n], | |
| 526: [21n, 9n, 2n, 0n], | |
| 706: [25n, 9n, 0n, 0n], | |
| 730: [27n, 1n, 0n, 0n], | |
| 829: [27n, 10n, 0n, 0n], | |
| 1414: [33n, 18n, 1n, 0n], | |
| 1549: [35n, 18n, 0n, 0n], | |
| 1906: [41n, 15n, 0n, 0n], | |
| 2986: [45n, 31n, 0n, 0n], | |
| 7549: [85n, 18n, 0n, 0n], | |
| 9634: [97n, 15n, 0n, 0n], | |
| }; | |
| function getBitWidth(n) { | |
| let e = 1n; | |
| while (1n << e <= n) { | |
| ++e; | |
| } | |
| return e - 1n; | |
| } | |
| function getMin(a, b) { | |
| return a > b ? b : a; | |
| } | |
| function getIntSqrt(n) { | |
| if (n <= 1n) { | |
| return n; | |
| } | |
| let x0 = n; | |
| let x1 = getMin(1n << (getBitWidth(n) / 2n + 1n), n / 2n); | |
| while (x1 < x0) { | |
| x0 = x1; | |
| x1 = (x0 + n / x0) / 2n; | |
| } | |
| return x0; | |
| } | |
| function isPerfectSquare(n) { | |
| const sqrt = getIntSqrt(n); | |
| return sqrt * sqrt === n; | |
| } | |
| function getGCD(a, b) { | |
| return b === 0n ? a : getGCD(b, a % b); | |
| } | |
| function getAbs(a) { | |
| return a > 0 ? a : -a; | |
| } | |
| function getModPow(val, exp, mod) { | |
| let result = 1n, | |
| subPower = val % mod; | |
| while (exp > 0n) { | |
| if ((exp & 1n) === 1n) { | |
| result = (result * subPower) % mod; | |
| } | |
| exp >>= 1n; | |
| subPower = (subPower * subPower) % mod; | |
| } | |
| return result; | |
| } | |
| function doMillerRabinTest(n, a) { | |
| let d = n - 1n; | |
| while (d % 2n === 0n) { | |
| d /= 2n; | |
| } | |
| let x = getModPow(a, d, n); | |
| if (x === 1n || x === n - 1n) { | |
| return true; | |
| } | |
| while (d !== n - 1n) { | |
| x = (x * x) % n; | |
| d *= 2n; | |
| if (x === n - 1n) { | |
| return true; | |
| } | |
| } | |
| return false; | |
| } | |
| class RNG { | |
| constructor() { | |
| this.seed = 1n; | |
| } | |
| next31Bit() { | |
| this.seed = (1103515245n * this.seed + 12345n) % 0x7fffffffn; | |
| return this.seed; | |
| } | |
| nextRange(bound) { | |
| const width = getBitWidth(bound); | |
| let result = 0n; | |
| for (let pos = 0n; pos < width; pos += 31n) { | |
| result = (result << 31n) | this.next31Bit(); | |
| } | |
| return result % bound; | |
| } | |
| } | |
| function isPrime(n) { | |
| if (n === 2n) { | |
| return true; | |
| } else if (n <= 1n || n % 2n === 0n) { | |
| return false; | |
| } | |
| let robustTests = ROBUST_TEST_SET.find(([bound]) => n < bound)?.[1]; | |
| if (robustTests === undefined) { | |
| const rng = new RNG(); | |
| robustTests = Array.from({ length: 15 }, () => rng.next31Bit()); | |
| } | |
| return robustTests.every((a) => { | |
| return doMillerRabinTest(n, a); | |
| }); | |
| } | |
| function factorizeRecursively(n, result) { | |
| if (n <= 1n) { | |
| return; | |
| } else if (isPrime(n)) { | |
| result.push(n); | |
| return; | |
| } else if (n % 2n === 0n) { | |
| result.push(2n); | |
| factorizeRecursively(n / 2n, result); | |
| return; | |
| } | |
| const rng = new RNG(); | |
| let xs, | |
| xt, | |
| c, | |
| factor = n; | |
| const f = (x) => (((x * x) % n) + c) % n; | |
| do { | |
| if (factor === n) { | |
| xs = xt = rng.nextRange(n - 2n) + 2n; | |
| c = rng.nextRange(20n) + 1n; | |
| } | |
| xs = f(xs); | |
| xt = f(f(xt)); | |
| factor = getGCD(getAbs(xs - xt), n); | |
| } while (factor === 1n || factor === n); | |
| factorizeRecursively(factor, result); | |
| factorizeRecursively(n / factor, result); | |
| } | |
| function factorize(n) { | |
| const result = []; | |
| factorizeRecursively(n, result); | |
| return result.sort((a, b) => (a < b ? -1 : a > b ? 1 : 0)); | |
| } | |
| function mergeAsTwo(a, b) { | |
| return [getAbs(a[0] * b[0] - a[1] * b[1]), a[0] * b[1] + a[1] * b[0], 0n, 0n]; | |
| } | |
| function getRemainders(a, b, bound) { | |
| const boundSqrt = getIntSqrt(bound); | |
| const result = []; | |
| while (b !== 0n) { | |
| const r = a % b; | |
| if (r <= boundSqrt) { | |
| result.push(r); | |
| } | |
| a = b; | |
| b = r; | |
| } | |
| return result; | |
| } | |
| function findSumOfTwo(p) { | |
| if (p === 1n) { | |
| return [1n, 0n, 0n, 0n]; | |
| } | |
| if (p === 2n) { | |
| return [1n, 1n, 0n, 0n]; | |
| } | |
| let rems = []; | |
| const rng = new RNG(); | |
| do { | |
| let q = rng.nextRange(p); | |
| while (getModPow(q, (p - 1n) / 2n, p) !== p - 1n) { | |
| q = rng.nextRange(p); | |
| } | |
| const x = getModPow(q, (p - 1n) / 4n, p); | |
| rems = getRemainders(p, x, p); | |
| } while (rems.length <= 2); | |
| return [rems[0], rems[1], 0n, 0n]; | |
| } | |
| function findSumOfThree(n) { | |
| if (RABIN_EXCEPTIONS[n]) { | |
| return RABIN_EXCEPTIONS[n]; | |
| } else if (n % 4n === 0n) { | |
| const sub = findSumOfThree(n / 4n); | |
| return [sub[0] * 2n, sub[1] * 2n, sub[2] * 2n, sub[3] * 2n]; | |
| } else if (n % 8n === 7n) { | |
| return [0n, 0n, 0n, 0n]; // Exception - Impossible | |
| } else if (n % 8n === 3n) { | |
| let x, p; | |
| const rng = new RNG(); | |
| do { | |
| x = rng.nextRange(getIntSqrt(n) + 1n); | |
| p = (n - x * x) / 2n; | |
| } while ((n - x * x) % 2n !== 0n || !(isPrime(p) || p === 1n)); | |
| const two = findSumOfTwo(p); | |
| return [x, two[0] + two[1], getAbs(two[0] - two[1]), 0n]; | |
| } else if (isPerfectSquare(n)) { | |
| return [getIntSqrt(n), 0n, 0n, 0n]; | |
| } else { | |
| let x, p; | |
| const rng = new RNG(); | |
| do { | |
| x = rng.nextRange(getIntSqrt(n) + 1n); | |
| p = n - x * x; | |
| } while (!isPrime(p)); | |
| const two = findSumOfTwo(p); | |
| return [x, two[0], two[1], 0n]; | |
| } | |
| } | |
| function fourSquares(n) { | |
| if (typeof n !== "bigint") { | |
| n = BigInt(n); | |
| } | |
| if (n === 1n) { | |
| return [1n, 0n, 0n, 0n]; | |
| } else if (n === 2n) { | |
| return [1n, 1n, 0n, 0n]; | |
| } else if (n === 3n) { | |
| return [1n, 1n, 1n, 0n]; | |
| } | |
| if (isPerfectSquare(n)) { | |
| return [getIntSqrt(n), 0n, 0n, 0n]; | |
| } | |
| if (n % 8n === 7n) { | |
| const sub = findSumOfThree(n - 4n); | |
| return [2n, sub[0], sub[1], sub[2]]; | |
| } | |
| return findSumOfThree(n); | |
| } |
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
https://www.codewars.com/kata/63022799acfb8d00285b4ea0