3 files changed, 49 insertions, 23 deletions

diff --git a/src/2016/day2/aoc.cpp b/src/2016/day2/aoc.cpp
index 3b5b28b..a73f26f 100644
--- a/src/2016/day2/aoc.cpp
+++ b/src/2016/day2/aoc.cpp
@@ -2,14 +2,16 @@
 
 namespace aoc2016 {
 
-int day2(line_view file) {
+int day2(line_view file, char ds[]) {
   int d{0};
-  int l{5};
+  size_t i{0};
+  char l{5};
   keybad pad;
-  per_line(file, [&d, &l, &pad](line_view lv) {
-    l = pad.touch(l, lv.line);
+  per_line(file, [&d, &l, &pad, &i, &ds](line_view lv) {
+    l = pad.touch(l, lv.line, pad.digit_keys);
     // printf("%d\n", l);
     d = d * 10 + l;
+    pad.touch(i++, ds, lv.line, pad.diamond_keys);
     return true;
   });
   return d;
diff --git a/src/2016/day2/aoc.h b/src/2016/day2/aoc.h
index 2b444cd..9a587d3 100644
--- a/src/2016/day2/aoc.h
+++ b/src/2016/day2/aoc.h
@@ -1,5 +1,7 @@
 #pragma once
 #include "common.h"
+#include <map>
+#include <vector>
 
 namespace aoc2016 {
 
@@ -11,33 +13,51 @@ struct keybad {
     left,
   };
 
-  int keys[9][4] = {
-      {1, 2, 4, 1}, {2, 3, 5, 1}, {3, 3, 6, 2}, {1, 5, 7, 4}, {2, 6, 8, 4},
-      {3, 6, 9, 5}, {4, 8, 7, 7}, {5, 9, 8, 7}, {6, 9, 9, 8},
+  std::map<char, std::vector<char>> digit_keys = {
+      {'1', {1, 2, 4, 1}}, {'2', {2, 3, 5, 1}}, {'3', {3, 3, 6, 2}}, {'4', {1, 5, 7, 4}}, {'5', {2, 6, 8, 4}},
+      {'6', {3, 6, 9, 5}}, {'7', {4, 8, 7, 7}}, {'8', {5, 9, 8, 7}}, {'9', {6, 9, 9, 8}},
   };
 
-  int touch(int i, char c) {
-    static int az[] = {
-        INT_MAX, INT_MAX, INT_MAX,
-        2,                                                                // D
-        INT_MAX, INT_MAX, INT_MAX, INT_MAX, INT_MAX, INT_MAX, INT_MAX, 3, // L
-        INT_MAX, INT_MAX, INT_MAX, INT_MAX, INT_MAX,
-        1,                   // R
-        INT_MAX, INT_MAX, 0, // U
-        INT_MAX, INT_MAX, INT_MAX, INT_MAX, INT_MAX,
-    };
-    return keys[i - 1][az[c - 'A']];
-  }
+  std::map<char, std::vector<char>> diamond_keys = {
+      {'1', {'1', '1', '3', '1'}}, {'2', {'2', '3', '6', '2'}}, {'3', {'1', '4', '7', '2'}},
+      {'4', {'4', '4', '8', '3'}}, {'5', {'5', '6', '5', '5'}}, {'6', {'2', '7', 'A', '5'}},
+      {'7', {'3', '8', 'B', '6'}}, {'8', {'4', '9', 'C', '7'}}, {'9', {'9', '9', '9', '8'}},
+      {'A', {'6', 'B', 'A', 'A'}}, {'B', {'7', 'C', 'D', 'A'}}, {'C', {'8', 'C', 'C', 'B'}},
+      {'D', {'B', 'D', 'D', 'D'}},
+  };
 
-  int touch(int i, const char* p) {
+  // clang-format off
+  int az[26] = {
+      INT_MAX, INT_MAX, INT_MAX, 2,
+      INT_MAX, INT_MAX, INT_MAX,
+      INT_MAX, INT_MAX,
+      INT_MAX, INT_MAX, 3,
+      INT_MAX, INT_MAX, INT_MAX,
+      INT_MAX, INT_MAX, 1,
+      INT_MAX, INT_MAX, 0,
+      INT_MAX, INT_MAX, INT_MAX,
+      INT_MAX, INT_MAX,
+  };
+  // clang-format on
+
+  char touch(char i, const char* p, std::map<char, std::vector<char>>& keys) {
     while (*p >= 'A' && *p <= 'Z') {
-      i = touch(i, *p);
+      i = keys[char(i + '0')][az[*p - 'A']];
       p++;
     }
     return i;
   }
+
+  void touch(size_t i, char codes[], const char* p, std::map<char, std::vector<char>>& keys) {
+    char x = i == 0 ? '5' : codes[i - 1];
+    while (*p >= 'A' && *p <= 'Z') {
+      x = keys[x][az[*p - 'A']];
+      p++;
+    }
+    codes[i] = x;
+  }
 };
 
-int day2(line_view);
+int day2(line_view, char[]);
 
 } // namespace aoc2016
diff --git a/test/test_2016.cpp b/test/test_2016.cpp
index 87b674f..288b8f7 100644
--- a/test/test_2016.cpp
+++ b/test/test_2016.cpp
@@ -2,6 +2,7 @@
 #include "2016/day2/aoc.h"
 #include "catch.hpp"
 #include <stdio.h>
+#include <string.h>
 
 TEST_CASE("No Time for a Taxicab", "[2016]") {
   line_view lv = load_file("../src/2016/day1/input");
@@ -12,5 +13,8 @@ TEST_CASE("No Time for a Taxicab", "[2016]") {
 
 TEST_CASE("Bathroom Security", "[2016]") {
   line_view lv = load_file("../src/2016/day2/input");
-  REQUIRE(82958 == aoc2016::day2(lv));
+  char codes[6] = {0};
+  int code = aoc2016::day2(lv, codes);
+  REQUIRE(82958 == code);
+  REQUIRE(strcmp(codes, "B3DB8") == 0);
 }