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

use crate::input::Input;
use crate::mod_exp::mod_exp;
use std::collections::HashMap;

const MODULO: u64 = 20_201_227;

/// Computes `loop_size` so that `BASE ^ loop_size % MODULO == public_key`.
///
/// See https://en.wikipedia.org/wiki/Baby-step_giant-step
fn babystep_giantstep(public_key: u32) -> u32 {
    // math.ceil(math.sqrt(MODULO-1)):
    const SQRT_MODULO_MINUS_ONE: u32 = 4_495;
    const BASE: u64 = 7;
    // mod_exp(BASE, SQRT_MODULO_MINUS_ONE * (MODULO - 2), MODULO):
    const FACTOR: u64 = 680_915;

    // Store the mapping from `BASE ^ x % MODULO` to `x` in a table:
    let baby_table = (0..SQRT_MODULO_MINUS_ONE)
        .scan(1_u32, |table_key, table_value| {
            let entry = Some((*table_key, table_value));
            *table_key = ((u64::from(*table_key) * BASE) % MODULO) as u32;
            entry
        })
        .collect::<HashMap<u32, u32>>();

    // https://stackoverflow.com/a/36893730/300710
    // Fermat little theorem states that
    //   "If p is a prime number, then for any integer a, the number a^p − a is an integer multiple of p."
    // Which is written as:
    //   a^p ≡ a (mod p)
    // If a is not divisible by p, it means that we can divide with a to get:
    //   a^(p-1) ≡ 1 (mod p)
    // or
    //   a * a^(p-2) ≡ 1 (mod p)
    // Which means that `a^(p-2)` is the multiplicative inverse of `a` (mod p), since the definition
    // of a multiplicative inverse `b` is: `a*b = 1 (mod p)`.
    //
    // So here:
    //    FACTOR := BASE^(SQRT_MODULO_MINUS_ONE * (MODULO - 2)) % MODULO
    // is the multiplicative inverse of `BASE` raised with SQRT_MODULO_MINUS_ONE,
    // so `BASE^-SQRT_MODULO_MINUS_ONE (mod p)`

    (0..SQRT_MODULO_MINUS_ONE)
        .try_fold(public_key, |state, giant_step| {
            baby_table.get(&state).map_or_else(
                // No entry - continue with new state `(state * FACTOR) % MODULO`:
                || Ok(((u64::from(state) * FACTOR) % MODULO) as u32),
                // We have found x in `BASE ^ x % MODULO = state`, where
                //   state = (public_key * FACTOR ^ giant_step) % MODULO.
                // Multiply with `SQRT_MODULO_MINUS_ONE` (of which is `FACTOR` is
                // the multiplicative inverse) `giant_step` times to get back `public_key`.
                |x| Err(giant_step * SQRT_MODULO_MINUS_ONE + x),
            )
        })
        .unwrap_err()
}

pub fn solve(input: &mut Input) -> Result<u64, String> {
    if input.is_part_two() {
        return Ok(0);
    }

    let on_error = || "Invalid input".to_string();

    let mut lines = input.text.lines();

    let card_public_key = lines
        .next()
        .ok_or_else(on_error)?
        .parse::<u32>()
        .map_err(|_| on_error())?;

    let door_public_key = lines
        .next()
        .ok_or_else(on_error)?
        .parse::<u32>()
        .map_err(|_| on_error())?;

    let card_loop_size = babystep_giantstep(card_public_key);

    let encryption_key = mod_exp(
        i128::from(door_public_key),
        i128::from(card_loop_size),
        i128::from(MODULO),
    ) as u64;

    Ok(encryption_key as u64)
}

#[test]
pub fn tests() {
    use crate::test_part_one;

    let example = "5764801\n17807724";
    test_part_one!(example => 14_897_079);
    let real_input = include_str!("day25_input.txt");
    test_part_one!(real_input => 18_862_163);
}