Advent of Code 2020 Day 18

advent_of_code_2020_day18.swift

indirect enum Element {
    case number(Int)
    case operation(Character)
    case subexpression([Element])
}

func parse(_ remainder: String.SubSequence, _ accumulator: [Element] = []) -> [Element] {
    guard let head = remainder.first else { return accumulator }
    switch head {
    case " ":
        return parse(remainder.dropFirst(), accumulator)
    case "0"..."9":
        guard let number = Int(remainder.prefix(while: \.isNumber)) else { fatalError("could not parse \(head)") }
        return parse(remainder.drop(while: \.isNumber), accumulator + [.number(number)])
    case "+", "*":
        return parse(remainder.dropFirst(), accumulator + [.operation(head)])
    case "(":
        let tail = remainder.dropFirst()
        let open = tail.enumerated().filter { $0.1 == "(" }.map(\.offset)
        let close = tail.enumerated().filter { $0.1 == ")" }.map(\.offset)
        guard let cutoff = zip(open, close).first(where: { open, close in close < open })?.1 ?? close.last else {
            fatalError("open/close failed: \(open) & \(close)")
        }
        let index = String.Index(utf16Offset: cutoff, in: tail)
        return parse(tail.dropFirst(cutoff + 1), accumulator + [.subexpression(parse(tail.prefix(upTo: index)))])
    default:
        fatalError("encountered unknown character: \(head)")
    }
}

// For part two. This will chunk anything on either side of a '*' into a subexpression
func preprocess(_ elements: [Element]) -> [Element] {
    let processed = elements.reduce(into: (acc: [Element](), sub: [Element]()), { acc, element in
        switch element {
        case .operation("*"):
            acc.acc.append(.subexpression(acc.sub))
            acc.acc.append(element)
            acc.sub = []
        case .subexpression(let els):
            acc.sub.append(.subexpression(preprocess(els)))
        default:
            acc.sub.append(element)
        }
    })
    return processed.acc + [.subexpression(processed.sub)]
}

typealias Op = (Int, Int) -> Int

func evaluate(_ elements: [Element]) -> Int {
    elements.reduce(into: (sum: 0, op: { _, s in s } as Op), { acc, element in
        switch element {
        case .number(let n):
            acc.sum = acc.sum == 0 ? n : acc.op(n, acc.sum)
        case .operation(let o):
            acc.op = { n, sum in
                switch o {
                case "+": return sum + n
                case "*": return max(sum, 1) * n
                default: fatalError()
                }
            }
        case .subexpression(let els):
            let evaluated = evaluate(els)
            acc.sum = acc.sum == 0 ? evaluated : acc.op(evaluated, acc.sum)
        }
    }).sum
}

// `input` is an array of strings

let day18a = input.map { parse($0[...]) }.map(evaluate).reduce(0, +)
print(day18a)
let day18b = input.map { parse($0[...]) }.map(preprocess).map(evaluate).reduce(0, +)
print(day18b)