From 6270ea51ea61ed8f9f5dd4999200d6212a31e268 Mon Sep 17 00:00:00 2001
From: Nick Krichevsky <nick@ollien.com>
Date: Mon, 14 Dec 2020 02:33:06 -0500
Subject: [PATCH] Add solution for day 14 part 2

---
 day14/day14.cpp | 68 ++++++++++++++++++++++++++++++++++++++++++++-----
 1 file changed, 61 insertions(+), 7 deletions(-)

diff --git a/day14/day14.cpp b/day14/day14.cpp
index 6562929..b228d46 100644
--- a/day14/day14.cpp
+++ b/day14/day14.cpp
@@ -8,6 +8,7 @@
 #include <map>
 #include <numeric>
 #include <regex>
+#include <set>
 #include <vector>
 
 constexpr char IGNORE_CHAR = 'X';
@@ -21,22 +22,24 @@ class InstructionBlock {
 	}
 
 	// We are dealing with a 36 bit numebr so we must use a long long to guarantee we can fit it
-	long long maskNumber(long long num) const {
+	long long maskValue(long long num) const {
 		long long res = num;
 		for (int i = 0; i < this->mask.size(); i++) {
 			char maskChar = this->mask.at(this->mask.size() - i - 1);
 			if (maskChar == IGNORE_CHAR) {
 				continue;
-			} else if (maskChar == '0') {
-				res &= ~(1LL << i);
-			} else if (maskChar == '1') {
-				res |= (1LL << i);
+			} else {
+				res = this->setBitAt(res, i, maskChar);
 			}
 		}
 
 		return res;
 	}
 
+	std::vector<long long> maskMemoryAddress(long long address) const {
+		return this->recursivelyMaskAddresses(address);
+	}
+
 	const std::vector<std::pair<int, int>> &getStoreInstructions() const {
 		return this->storeInstructions;
 	}
@@ -44,6 +47,40 @@ class InstructionBlock {
  private:
 	std::string mask;
 	std::vector<std::pair<int, int>> storeInstructions;
+
+	long long setBitAt(long long value, int position, char bit) const {
+		auto res = value;
+		if (bit == '0') {
+			res &= ~(1LL << position);
+		} else if (bit == '1') {
+			res |= (1LL << position);
+		} else {
+			throw std::invalid_argument("bit must be zero or one");
+		}
+
+		return res;
+	}
+
+	std::vector<long long> recursivelyMaskAddresses(long long address, int startIdx = 0) const {
+		std::vector<long long> results;
+		long long res = address;
+		for (int i = startIdx; i < this->mask.size(); i++) {
+			char maskChar = this->mask.at(this->mask.size() - i - 1);
+			if (maskChar == IGNORE_CHAR) {
+				auto masked0 = this->setBitAt(res, i, '0');
+				auto masked1 = this->setBitAt(res, i, '1');
+				auto results0 = this->recursivelyMaskAddresses(masked0, i + 1);
+				auto results1 = this->recursivelyMaskAddresses(masked1, i + 1);
+				results.insert(results.end(), results0.begin(), results0.end());
+				results.insert(results.end(), results1.begin(), results1.end());
+			} else if (maskChar == '1') {
+				res = this->setBitAt(res, i, maskChar);
+			}
+		}
+
+		results.push_back(res);
+		return results;
+	}
 };
 
 std::vector<std::string> readInput(const std::string &filename) {
@@ -91,8 +128,24 @@ long long part1(const std::vector<InstructionBlock> &instructions) {
 	std::unordered_map<int, long long> memory;
 	for (const InstructionBlock &instruction : instructions) {
 		for (const std::pair<int, int> &storeInstruction : instruction.getStoreInstructions()) {
-			auto maskedStorage = instruction.maskNumber(storeInstruction.second);
-			memory[storeInstruction.first] = maskedStorage;
+			auto maskedValue = instruction.maskValue(storeInstruction.second);
+			memory[storeInstruction.first] = maskedValue;
+		}
+	}
+
+	return std::accumulate(
+		memory.cbegin(), memory.cend(), 0LL, [](long long total, std::pair<int, long long> memoryItem) {
+			return total + memoryItem.second;
+		});
+}
+
+long long part2(const std::vector<InstructionBlock> &instructions) {
+	std::unordered_map<long long, long long> memory;
+	for (const InstructionBlock &instruction : instructions) {
+		for (const std::pair<int, int> &storeInstruction : instruction.getStoreInstructions()) {
+			for (const auto maskedAddress : instruction.maskMemoryAddress(storeInstruction.first)) {
+				memory[maskedAddress] = storeInstruction.second;
+			}
 		}
 	}
 
@@ -112,4 +165,5 @@ int main(int argc, char *argv[]) {
 	auto parsedInput = parseInput(input);
 
 	std::cout << part1(parsedInput) << std::endl;
+	std::cout << part2(parsedInput) << std::endl;
 }