1  #![feature(test)]
 2  
 3  #[derive(Debug)]
 4  struct Input {
 5      races: Vec<Race>,
 6      part2: Race,
 7  }
 8  
 9  #[derive(Debug, Clone, Copy)]
10  struct Race {
11      time: u64,
12      distance: u64,
13  }
14  
15  fn setup(input: &str) -> Input {
16      let mut lines = input.lines();
17      let mut parse_line = || {
18          lines
19              .next()
20              .unwrap()
21              .split_whitespace()
22              .skip(1)
23              .map(|x| x.parse().unwrap())
24      };
25      let races = parse_line()
26          .zip(parse_line())
27          .map(|(t, d)| Race {
28              time: t,
29              distance: d,
30          })
31          .collect();
32  
33      let mut lines = input.lines();
34      let mut parse_line = || {
35          lines
36              .next()
37              .unwrap()
38              .split(':')
39              .nth(1)
40              .unwrap()
41              .replace(' ', "")
42              .parse()
43              .unwrap()
44      };
45      let part2 = Race {
46          time: parse_line(),
47          distance: parse_line(),
48      };
49  
50      Input { races, part2 }
51  }
52  
53  fn solve(Race { time, distance }: Race) -> u64 {
54      let Some(x) = (time * time).checked_sub(distance * 4) else {
55          return 0;
56      };
57      let root = x.isqrt();
58      let root_is_int = (root * root == x) as u64;
59      root + (time + root + 1 + root_is_int) % 2 - root_is_int
60  }
61  
62  fn part1(input: &Input) -> u64 {
63      input
64          .races
65          .iter()
66          .map(|&race| solve(race))
67          .filter(|&x| x > 0)
68          .product()
69  }
70  
71  fn part2(input: &Input) -> u64 {
72      solve(input.part2)
73  }
74  
75  aoc::main!(2023, 6, ex: 1);
76  
77  #[cfg(test)]
78  mod tests {
79      use proptest::{prop_assert_eq, proptest};
80  
81      use super::*;
82  
83      proptest! {
84          #[test]
85          fn proptest_solve(t in (0u64..=200000), d in (0u64..=1000000)) {
86              let race = Race { time: t, distance: d };
87              let expected = (0..=t).filter(|i| i * (t - i) > d).count() as u64;
88  
89              let result = solve(race);
90  
91              prop_assert_eq!(result, expected);
92          }
93      }
94  }