main.rs (5700B)
1 use std::collections::{HashMap, HashSet, VecDeque};
2
3 #[derive(Debug, Copy, Clone)]
4 enum CaveSize {
5 Big,
6 Small,
7 }
8
9 impl CaveSize {
10 fn check(s: &str) -> CaveSize {
11 if s.chars().all(char::is_lowercase) {
12 CaveSize::Small
13 } else {
14 CaveSize::Big
15 }
16 }
17 }
18
19 #[derive(Debug, Clone)]
20 struct CaveSystem {
21 map: HashMap<String, HashSet<String>>,
22 cave_size: HashMap<String, CaveSize>,
23 }
24
25 impl CaveSystem {
26 fn new(input: &Vec<String>) -> Self {
27 let mut map = HashMap::<String, HashSet<String>>::new();
28 let mut cave_size = HashMap::<String, CaveSize>::new();
29 for s in input {
30 let mut tmp = s.split('-');
31 let a = tmp.next().unwrap().to_string();
32 let b = tmp.next().unwrap().to_string();
33 if let Some(connected) = map.get_mut(&a) {
34 connected.insert(b.clone());
35 } else {
36 map.insert(a.clone(), {
37 let mut tmp = HashSet::new();
38 tmp.insert(b.clone());
39 tmp
40 });
41 }
42 if !cave_size.contains_key(&a) {
43 cave_size.insert(a.clone(), CaveSize::check(&a));
44 }
45 if let Some(connected) = map.get_mut(&b) {
46 connected.insert(a.clone());
47 } else {
48 map.insert(b.clone(), {
49 let mut tmp = HashSet::new();
50 tmp.insert(a.clone());
51 tmp
52 });
53 }
54 if !cave_size.contains_key(&b) {
55 cave_size.insert(b.clone(), CaveSize::check(&b));
56 }
57 }
58 CaveSystem { map, cave_size }
59 }
60 }
61
62 fn main() {
63 let input = std::fs::read_to_string("input.txt")
64 .unwrap()
65 .trim()
66 .split('\n')
67 .map(|s| s.to_string())
68 .collect::<Vec<_>>();
69 let cave_system = CaveSystem::new(&input);
70 // 5076
71 println!("Part 1: {}", part_1(&cave_system));
72 // 145643
73 println!("Part 2: {}", part_2(&cave_system));
74 }
75
76 fn part_1(cave_system: &CaveSystem) -> i32 {
77 let cave_system = cave_system.clone();
78 let mut queue = VecDeque::<(VecDeque<String>, HashSet<String>)>::new();
79 queue.push_back({
80 let mut path = VecDeque::<String>::new();
81 path.push_back("start".to_string());
82 (path, HashSet::<String>::new())
83 });
84 let mut path_count = 0;
85 while let Some((path, mut visited)) = queue.pop_front() {
86 if let Some(last_step) = path.back() {
87 visited.insert(last_step.clone());
88 if last_step == "end" {
89 path_count += 1;
90 } else {
91 for next_step in cave_system.map[last_step].iter() {
92 match (
93 cave_system.cave_size[next_step],
94 visited.contains(next_step),
95 ) {
96 (CaveSize::Big, _) | (CaveSize::Small, false) => {
97 queue.push_back((
98 {
99 let mut tmp = path.clone();
100 tmp.push_back(next_step.clone());
101 tmp
102 },
103 visited.clone(),
104 ));
105 }
106 _ => (),
107 }
108 }
109 }
110 }
111 }
112 path_count
113 }
114
115 fn part_2(cave_system: &CaveSystem) -> i32 {
116 let cave_system = cave_system.clone();
117 let mut queue = VecDeque::<(VecDeque<String>, HashSet<String>, Option<String>)>::new();
118 queue.push_back({
119 let mut path = VecDeque::<String>::new();
120 path.push_back("start".to_string());
121 (path, HashSet::<String>::new(), None)
122 });
123 let mut path_count = 0;
124 while let Some((path, mut visited, twice)) = queue.pop_front() {
125 if let Some(last_step) = path.back() {
126 visited.insert(last_step.clone());
127 if last_step == "end" {
128 path_count += 1;
129 } else {
130 for next_step in cave_system.map[last_step].iter() {
131 match (
132 cave_system.cave_size[next_step],
133 visited.contains(next_step),
134 twice.clone(),
135 ) {
136 (CaveSize::Big, _, twice) | (CaveSize::Small, false, twice) => {
137 queue.push_back((
138 {
139 let mut tmp = path.clone();
140 tmp.push_back(next_step.clone());
141 tmp
142 },
143 visited.clone(),
144 twice,
145 ));
146 }
147 (CaveSize::Small, true, None) => {
148 if next_step != "start" {
149 queue.push_back((
150 {
151 let mut tmp = path.clone();
152 tmp.push_back(next_step.clone());
153 tmp
154 },
155 visited.clone(),
156 Some(next_step.clone()),
157 ));
158 }
159 }
160 _ => (),
161 }
162 }
163 }
164 }
165 }
166 path_count
167 }