common.rs (5314B)
1 // This works for part 1, but it is too slow for part 2. 2 // May be I am doing something wrong with caching? 3 4 use std::cmp::{max, Ordering}; 5 use std::ops::{AddAssign, SubAssign}; 6 use std::collections::HashMap; 7 8 #[derive(Debug, Copy, Clone, PartialEq, Eq, Hash)] 9 pub struct Stock { 10 pub ore: i32, 11 pub cla: i32, 12 pub obs: i32 13 } 14 15 impl Stock { 16 pub fn all_lower(&self, other: &Self) -> bool { 17 self.ore <= other.ore && self.cla <= other.cla && self.obs <= other.obs 18 } 19 } 20 21 impl PartialOrd for Stock { 22 fn partial_cmp(&self, other: &Self) -> Option<Ordering> { 23 if self == other { return Some(Ordering::Equal); } 24 if self.all_lower(other) { return Some(Ordering::Less); } 25 if other.all_lower(self) { return Some(Ordering::Greater); } 26 None 27 } 28 } 29 30 impl AddAssign<&Stock> for Stock { 31 fn add_assign(&mut self, other: &Self) { 32 self.ore += other.ore; 33 self.cla += other.cla; 34 self.obs += other.obs; 35 } 36 } 37 38 impl SubAssign<&Stock> for Stock { 39 fn sub_assign(&mut self, other: &Self) { 40 self.ore -= other.ore; 41 self.cla -= other.cla; 42 self.obs -= other.obs; 43 } 44 } 45 46 #[derive(Debug, Copy, Clone, PartialEq, Eq, Hash)] 47 pub struct Status { 48 pub bots: Stock, 49 pub resources: Stock 50 } 51 52 impl Status { 53 pub fn new() -> Status { 54 Status { 55 bots: Stock { ore: 1, cla: 0, obs: 0 }, 56 resources: Stock { ore: 0, cla: 0, obs: 0 } 57 } 58 } 59 } 60 61 #[derive(Debug, Clone, Hash)] 62 pub struct Blueprint { 63 ore_bot_cost: Stock, 64 cla_bot_cost: Stock, 65 obs_bot_cost: Stock, 66 geo_bot_cost: Stock, 67 } 68 69 impl Blueprint { 70 fn parse_one_resource(line: &str, i: &mut usize) -> i32 { 71 *i = line[*i..].find("costs ").unwrap() + 6 + *i; 72 let j = line[*i..].find(' ').unwrap() + *i; 73 line[*i..j].parse::<i32>().unwrap() 74 } 75 76 fn parse_two_resources(line: &str, i: &mut usize) -> (i32, i32) { 77 *i = line[*i..].find("costs ").unwrap() + 6 + *i; 78 let j = line[*i..].find(' ').unwrap() + *i; 79 let a = line[*i..j].parse::<i32>().unwrap(); 80 81 *i = line[*i..].find("and ").unwrap() + 4 + *i; 82 let j = line[*i..].find(' ').unwrap() + *i; 83 let b = line[*i..j].parse::<i32>().unwrap(); 84 85 (a, b) 86 } 87 88 pub fn from_line(line: &str) -> Blueprint { 89 let mut i = 0; 90 91 let ore = Blueprint::parse_one_resource(&line, &mut i); 92 let ore_bot_cost = Stock { ore, cla: 0, obs: 0 }; 93 94 let ore = Blueprint::parse_one_resource(&line, &mut i); 95 let cla_bot_cost = Stock { ore, cla: 0, obs: 0 }; 96 97 let (ore, cla) = Blueprint::parse_two_resources(&line, &mut i); 98 let obs_bot_cost = Stock { ore, cla, obs: 0 }; 99 100 let (ore, obs) = Blueprint::parse_two_resources(&line, &mut i); 101 let geo_bot_cost = Stock { ore, cla: 0, obs }; 102 103 Blueprint { ore_bot_cost, cla_bot_cost, obs_bot_cost, geo_bot_cost } 104 } 105 } 106 107 pub fn read_blueprints_from_stdin() -> Vec<Blueprint> { 108 let mut v = vec![]; 109 let mut line = String::new(); 110 while std::io::stdin().read_line(&mut line).unwrap() > 0 { 111 v.push(Blueprint::from_line(&line)); 112 line.clear(); 113 } 114 v 115 } 116 117 pub fn most_geodes( 118 bp: &Blueprint, 119 s: &Status, 120 m: i32, 121 mem: &mut HashMap<(Status, i32), i64> 122 ) -> i64 { 123 if m <= 1 { return 0; } 124 125 if let Some(r) = mem.get(&(s.clone(), m)) { return *r; } 126 127 let mut new_status = s.clone(); 128 new_status.resources += &new_status.bots; 129 130 // If a geode bot can be built, it is always the best thing to do 131 if m > 1 && s.resources >= bp.geo_bot_cost { 132 new_status.resources -= &bp.geo_bot_cost; 133 let result = (m - 1) as i64 + most_geodes(bp, &new_status, m-1, mem); 134 mem.insert((s.clone(), m), result); 135 return result; 136 } 137 138 let mut result = 0; 139 140 // The second check is to avoid producing too many bots for each given 141 // type of resource. For example, 5 ore bots will produce more ore 142 // than the factory can ever consume, since all bots cost at most 4 ore. 143 let can_obs = s.resources >= bp.obs_bot_cost && 144 s.bots.obs < bp.geo_bot_cost.obs; 145 let can_cla = s.resources >= bp.cla_bot_cost && 146 s.bots.cla < bp.obs_bot_cost.cla; 147 let can_ore = s.resources >= bp.ore_bot_cost && 148 s.bots.ore < 4; 149 150 if m > 2 && can_obs { 151 new_status.resources -= &bp.obs_bot_cost; 152 new_status.bots.obs += 1; 153 result = max(result, most_geodes(bp, &new_status, m-1, mem)); 154 new_status.bots.obs -= 1; 155 new_status.resources += &bp.obs_bot_cost; 156 } 157 158 if m > 3 && can_cla { 159 new_status.resources -= &bp.cla_bot_cost; 160 new_status.bots.cla += 1; 161 result = max(result, most_geodes(bp, &new_status, m-1, mem)); 162 new_status.bots.cla -= 1; 163 new_status.resources += &bp.cla_bot_cost; 164 } 165 166 if m > 4 && can_ore { 167 new_status.resources -= &bp.ore_bot_cost; 168 new_status.bots.ore += 1; 169 result = max(result, most_geodes(bp, &new_status, m-1, mem)); 170 new_status.bots.ore -= 1; 171 new_status.resources += &bp.ore_bot_cost; 172 } 173 174 if !can_obs || !can_cla || !can_ore { 175 result = max(result, most_geodes(bp, &new_status, m-1, mem)); 176 } 177 178 mem.insert((s.clone(), m), result); 179 result 180 }