AoC 2020 - Day 14

day14.swift

import Foundation

extension String {
    var asUInt64: UInt64 { UInt64(self, radix: 2)! }

    func masking(with mask: String, mapper: (Character, Character) -> Character) -> [Element] {
        zip(padded(with: "0", toLength: 36), mask).map(mapper)
    }

    func padded(with character: String, toLength length: Int) -> String {
        return String(repeating: character, count: max(0, length - count)) + self
    }
}

extension UInt64 {
    var asBinaryString: String { String(self, radix: 2) }
}

extension Array where Element == String.Element {
    var string: String { String(self) }
}

enum Command {
    case setMask(String)
    case write(memory: UInt64, value: UInt64)

    init(string: String) {
        if string.starts(with: "mask") {
            let mask = string.components(separatedBy: "=")
                .map { $0.trimmingCharacters(in: .whitespaces) }[1]
            self = .setMask(mask)
        } else {
            let memAndValue = string.components(separatedBy: "=")
                .map { $0.trimmingCharacters(in: .whitespaces) }
            let value = UInt64(memAndValue[1])!
            let address = UInt64(memAndValue[0].trimmingCharacters(in: CharacterSet.decimalDigits.inverted))!
            self = .write(memory: address, value: value)
        }
    }
}

final class Day1420: Day {
    override func perform() {
        let commands = String.input(forDay: 14, year: 2020)
            .components(separatedBy: "\n")
            .map(Command.init)

        // part 1
        var memory: [UInt64: UInt64] = [:]
        var mask: String = ""
        for command in commands {
            switch command {
            case .setMask(let m): mask = m
            case let .write(address, value): memory[address] = masking(value: value, with: mask)
            }
        }

        stageOneOutput = "\(memory.values.reduce(0, +))"

        // part 2
        memory = [:]
        mask = ""
        for command in commands {
            switch command {
            case .setMask(let m):
                mask = m
            case let .write(address, value):
                maskedAddresses(address: address, with: mask).forEach { memory[$0] = value }
            }
        }

        stageTwoOutput = "\(memory.values.reduce(0, +))"
    }

    func masking(value: UInt64, with mask: String) -> UInt64? {
        value.asBinaryString
            .masking(with: mask) { (source, mask) in mask == "X" ? source : mask }
            .string
            .asUInt64
    }

    func maskedAddresses(address: UInt64, with mask: String) -> [UInt64] {
        let maskedAddress = address.asBinaryString.masking(with: mask) { (source, mask) in
            mask == "0" ? source : (mask == "1" ? mask : "X")
        }

        let xIndices = zip(maskedAddress, maskedAddress.indices)
            .filter { $0.0 == "X" }
            .reduce(into: []) { $0.append($1.1) }

        if xIndices.isEmpty { return [address] }

        return (0...String(repeating: "1", count: xIndices.count).asUInt64)
            .reduce(into: []) { (result, i) in
                let bits = i.asBinaryString.padded(with: "0", toLength: xIndices.count)
                var newResult = maskedAddress
                zip(bits, xIndices).forEach { bit, index in newResult[index] = bit }
                result.append(newResult.string.asUInt64)
            }
    }
}