1  #![feature(test)]
 2  
 3  use std::{
 4      cmp::Reverse,
 5      collections::{BinaryHeap, VecDeque},
 6  };
 7  
 8  type Input = Vec<u8>;
 9  
10  fn setup(input: &str) -> Input {
11      input.trim().bytes().map(|b| b - b'0').collect()
12  }
13  
14  fn part1(input: &Input) -> u64 {
15      let mut files = Vec::new();
16      let mut free = VecDeque::new();
17      let mut pos = 0;
18      for (i, len) in input.iter().map(|&x| x as u64).enumerate() {
19          if i & 1 == 0 {
20              files.push((pos, len, i as u64 >> 1));
21          } else if len > 0 {
22              free.push_back((pos, len));
23          }
24          pos += len;
25      }
26  
27      files
28          .iter()
29          .flat_map(|&(file_pos, file_len, file_id)| {
30              (file_pos..file_pos + file_len).map(move |pos| (pos, file_id))
31          })
32          .rev()
33          .map(|(mut pos, file_id)| {
34              if let Some((free_pos, free_len)) = free.front_mut().filter(|free| free.0 < pos) {
35                  pos = *free_pos;
36                  *free_pos += 1;
37                  *free_len -= 1;
38                  if *free_len == 0 {
39                      free.pop_front();
40                  }
41              }
42              pos * file_id
43          })
44          .sum()
45  }
46  
47  fn part2(input: &Input) -> u64 {
48      let mut files = Vec::new();
49      let mut free_by_len = [const { BinaryHeap::new() }; 10];
50      let mut pos = 0;
51      for (i, len) in input.iter().map(|&x| x as u64).enumerate() {
52          if i & 1 == 0 {
53              files.push((pos, len, i as u64 >> 1));
54          } else if len > 0 {
55              free_by_len[len as usize].push(Reverse(pos));
56          }
57          pos += len;
58      }
59  
60      files
61          .iter()
62          .rev()
63          .map(|&(file_pos, file_len, file_id)| {
64              let pos = (file_len as _..free_by_len.len())
65                  .filter_map(|len| free_by_len[len].peek().map(|&free_pos| (free_pos.0, len)))
66                  .min()
67                  .filter(|&(free_pos, _)| free_pos < file_pos)
68                  .map(|(free_pos, free_len)| {
69                      free_by_len[free_len].pop();
70                      if free_len > file_len as usize {
71                          free_by_len[free_len - file_len as usize]
72                              .push(Reverse(free_pos + file_len));
73                      }
74                      free_pos
75                  })
76                  .unwrap_or(file_pos);
77              file_len * (pos * 2 + file_len - 1) / 2 * file_id
78          })
79          .sum()
80  }
81  
82  aoc::main!(2024, 9, ex: 1);