1 files changed, 24 insertions, 13 deletions

diff --git a/01/src/main.rs b/01/src/main.rs
index a46ee99..21c164b 100644
--- a/01/src/main.rs
+++ b/01/src/main.rs
@@ -1,5 +1,7 @@
 use advent_lib::prelude::*;
 
+use std::collections::BTreeSet;
+
 
 fn main() -> Result<()> {
   let mut args = std::env::args();
@@ -10,25 +12,34 @@ fn main() -> Result<()> {
   let filename = args.next().unwrap();
 
   let mut input = advent_lib::read_int_file(&filename)?;
+
   input.sort();
 
+  let mut input_set = BTreeSet::new();
+  for item in &input {
+    input_set.insert(item);
+  }
+
   for i in 0 .. input.len() {
     let a = input[i];
     if a > 2020 {
       break;
     }
 
-    for j in i+1 .. input.len() {
-      let b = input[j];
-
-      if a + b == 2020 {
-        let product = a * b;
-        println!("a: {:?}, b: {:?}, a*b: {:?}", a, b, product);
-      }
+    let b = 2020 - a;
+    if input_set.contains(&b) {
+      let product = a * b;
+      println!("a: {:?}, b: {:?}, a*b: {:?}", a, b, product);
+      break;
     }
   }
 
+  let mut done = false;
   for i in 0 .. input.len() {
+    if done {
+      break;
+    }
+
     let a = input[i];
     if a > 2020 {
       break;
@@ -41,13 +52,13 @@ fn main() -> Result<()> {
         break;
       }
 
-      for k in j+1 .. input.len() {
-        let c = input[k];
+      let c = 2020 - a - b;
+      if input_set.contains(&c) {
+        let product = a * b * c;
+        println!("a: {:?}, b: {:?}, c: {:?}, a*b*c: {:?}", a, b, c, product);
 
-        if a + b + c == 2020 {
-          let product = a * b * c;
-          println!("a: {:?}, b: {:?}, c: {:?}, a*b*c: {:?}", a, b, c, product);
-        }
+        done = true;
+        break;
       }
     }
   }