Swift recursive enumerations and snailfish numbers (AoC 2021)
Sat, 08 Jan 2022 22:15:57 +0000
On day 18 of Advent of Code 2021 there was an interesting problem about something the author called "snailfish numbers". I immediately thought that, since I was solving these problems in Swift, a recursive enum would be ideal to represent these numbers. But the fanciest data structure is not always the best solution... 😅 Let's see why.

Snailfish numbers

The problem basically describes a snailfish number as a binary tree, where leafs of the tree are integer numbers. These trees are written down as lists of 2 elements, where each element can either be a number, or another list of two elements. Examples given in AoC:

[1,2]
[[1,2],3]
[9,[8,7]]
[[1,9],[8,5]]
[[[[1,2],[3,4]],[[5,6],[7,8]]],9]
[[[9,[3,8]],[[0,9],6]],[[[3,7],[4,9]],3]]
[[[[1,3],[5,3]],[[1,3],[8,7]]],[[[4,9],[6,9]],[[8,2],[7,3]]]]

Recursive enumerations

Enumerations in Swift can have associated values, and these enumerations can be defined recursively (you need to add the indirect keyword to the definition of your enum). I thought this would be a very neat declaration of a snailfish number (I call it SnailNumber, a bit shorter):
indirect enum SnailNumber {
    case number(Int)
    case pair(SnailNumber, SnailNumber)
}
It's pretty straightforward. This defines a binary tree. To understand how to use it, let's implement the CustomStringConvertible protocol to provide a textual description of a SnailNumber,
indirect enum SnailNumber: CustomStringConvertible {
    var description: String {
        get {
            switch(self){
            case .number(let number):
                return "\(number)"
            case .pair(let left, let right):
                return "[\(left),\(right)]"
            }
        }
    }
  
    case number(Int)
    case pair(SnailNumber, SnailNumber)
}
If you pay attention, that description is recursive, because "[\(left),\(right)]" invokes the description of left and the description of right, both SnailNumbers. It will stop when you find an integer number. Some example output:
let a = SnailNumber.pair(.number(1), .pair(.number(4), .number(5)))
print("\(a)")
// produces: [1,[4,5]]
The problem defines addition of snailfish numbers as simply the creation of a new pair, which it's very easy with these enums:
func + (lhs: SnailNumber, rhs: SnailNumber) -> SnailNumber {
    return .pair(lhs, rhs)
}  
Parsing a string with a snailfish number is a bit trickier. I'll paste here my code for reference:
    static func parse(_ s: String) -> SnailNumber {
      var rest = s
      var numbers: [SnailNumber] = []
      while !rest.isEmpty {
          if rest.first == "[" {
              // new pair
              let pair = SnailNumber.pair(.number(0), .number(0))
              numbers.append(pair)
              rest = rest.substring(from: 1)
          } else if rest.first == "]" {
              let right = numbers.removeLast()
              let left = numbers.removeLast()
              rest = rest.substring(from: 1)
              numbers[numbers.count - 1] = .pair(left, right)
          } else if rest.first == "," {
              rest = rest.substring(from: 1)
          } else {
              let groups = getGroupsFromRegex(pattern: #"^(\d+)(.*)"#, in: rest)
              if let pair = groups.first {
                  let n = SnailNumber.number(Int(pair[0])!)
                  numbers.append(n)
                  rest = pair[1]
              } else {
                  print("parse error: \(rest)")
              }

          }
      }
      return numbers.first!
  }

Exploding nightmares

Let's not celebrate yet. If you've read the AoC website, these numbers need to be reduced after every addition. The reduction process is an iterative process of explosions and splits. Read AoC2021 day 18 for details.

The splitting process keeps all the integer numbers smaller than 10, by splitting every number which it's 10 or greater into a new pair. For instance, 10 becomes [5,5], and 11 becomes [5,6]. Although this operation is easy with the enum I've defined, the problem says you need to apply one operation at a time, so I can't recursively split all the tree. It needs to be explored from left to right, and stop at the first split. I used Swift optionals (the ?) to check whether a split has happened, like this:

    static func split(_ n: SnailNumber) -> SnailNumber? {
      switch(n){
      case .number(let a):
          if a >= 10 {
              let r = a % 2
              return .pair(.number(a/2), .number((a+r)/2))
          }
      case .pair(let left, let right):
          if let sLeft = split(left) {
              // one split at a time at most
              return .pair(sLeft, right)
          }
          if let sRight = split(right) {
              return .pair(left, sRight)
          }
      }
      return nil
  }
The explosion operator is trickier. What it does is keeping the depth of the tree at 4 at maximum. If you find any node that is deeper, you take its left number and add it to the first regular number to its left, and its right number is added to the first regular number to its right. Again, you do such operations one at a time. When you look at snailfish numbers as strings, that sounds trivial. But with this recursive structure, things become complicated. Also, note that I can't modify a number in place, since enums are passed by value. I tried something similar to what I did for the split operator, but it ended being a nightmare. Here's the whole function for reference. Feel free to skip to the next paragraph.
  static func explode(_ n: SnailNumber, depth: Int) -> (n: SnailNumber?, left: Int?, right: Int?) {
    switch(n){
    case .number(_):
        return (nil, nil, nil)
    case .pair(let left, let right):
        switch(left, right) {
        case (let .number(a), let .number(b)):
            if depth >= 4 {
                return (.number(0), a, b)
            }
            return (nil, nil, nil)
        case (.pair(_, _), let .number(c)):
            let e = explode(left, depth: depth + 1)
            if let en = e.n {
                if let r = e.right {
                    let newPair = SnailNumber.pair(en, .number(c + r))
                    return (newPair, e.left, nil)
                } else {
                    return (SnailNumber.pair(en, .number(c)), e.left, e.right)
                }
            }
            return (nil, nil, nil)
        case (let .number(a), .pair(_, _)):
            let e = explode(right, depth: depth + 1)
            if let en = e.n {
                if let l = e.left {
                    let newPair = SnailNumber.pair(.number(a + l), en)
                    return (newPair, nil, e.right)
                } else {
                    return (SnailNumber.pair(.number(a), en), e.left, e.right)
                }
            }
            return (nil, nil, nil)
        case (.pair(_, _), .pair(_, _)):
            let eLeft = explode(left, depth: depth + 1)
            var leftPair: SnailNumber?
            var rightPair: SnailNumber?
            var propagateLeft: Int?
            var propagateRight: Int?
            if let en = eLeft.n {
                leftPair = en
                rightPair = right
                if let er = eLeft.right {
                    rightPair = SnailNumber.addLeft(right, value: er)
                }
                propagateLeft = eLeft.left
            }
            let eRight = explode(rightPair ?? right, depth: depth + 1)
            if let en = eRight.n {
                rightPair = en
                if leftPair == nil {
                    leftPair = left
                }
                if let el = eRight.left {
                    leftPair = SnailNumber.addLeft(leftPair!, value: el)
                }
                propagateRight = eRight.right
            }
            if let el = leftPair, let er = rightPair {
                return (.pair(el, er), propagateLeft, propagateRight)
            }
            return (nil, nil, nil)
        }
    }
}
This recursive function returns optional values that I use to know whether I need to propagate a number to the left or to the right. Because the tree is explored depth-first and from left to right, the numbers can be propagated to the right correctly, but they are only propagated correctly to the left if the number it's in an immediate parent of the exploded leaf. The algorithm worked for the examples given in the AoC website, but it didn't work for the general case.

While debugging this I realized this data structure looked cool on paper, but it had complicated things unnecessarily. 😣

Turing machines and flat snails

Compared to that binary tree, finding the left number in string format sounds more straightforward: simply move the cursor to the left until you find a number. If you've studied Turing machines at uni, that would probably sound familiar. So what if I stored the snailfish number as a string?

I decided that I would do some minimal parsing instead, and store it as a list of strings. Each string can either be a number, or a square bracket or a comma. The string description of such a number is just a join of all the elements in the list. Here's my flat snailfish number:

struct FlatSnail: CustomStringConvertible, Equatable {
    static let zero = FlatSnail("[0,0]")
    let number: [String]
    
    var description: String {
        get {
            return number.joined()
        }
    }
        
    init(number: [String]) {
        self.number = number
    }
    
    init(_ s: String) {
        var data: [String] = []
        var rest = s
        while !rest.isEmpty {
            if rest.first == "[" || rest.first == "]" || rest.first == "," {
                data.append(String(rest.first!))
                rest = rest.substring(from: 1)
            } else {
                // getGroupsFromRegex returns all the regex matches in [[String]]
                let groups = getGroupsFromRegex(pattern: #"^(\d+)(.*)"#, in: rest)
                if let pair = groups.first {
                    data.append(pair[0])
                    rest = pair[1]
                } else {
                    print("parse error: \(rest)")
                }

            }
        }
        number = data
    }
}
Addition with this structure is not as cool as with enums, but still quite simple:
func + (lhs: FlatSnail, rhs: FlatSnail) -> FlatSnail {
    let number = ["["] + lhs.number + [","] + rhs.number + ["]"]
    let snail = FlatSnail(number: number)
    return snail.reduce()
}
Notice that I'm already calling the reduce function because the problem specifies that numbers always need to be reduced after addition. The reduction process consists of iterating through explosions and splits until the number doesn't change. I wrote it like this:
func reduce() -> FlatSnail {
    var out = self
    var hasChanged = true
    while hasChanged {
        if let next = out.explode() {
            out = next
            continue
        }
        if let next = out.split() {
            out = next
            continue
        }
        hasChanged = false
    }
    return out
}  
And the explosion operator, although still a bit long, is now easier to implement because it's just moving the "cursor" from left to right. Every time you encounter an open square bracket, you are going down one level. A closing bracket means you go back one level. Here's the whole function:
func explode() -> FlatSnail? {
  var out: [String] = []
  var depth = 0
  for i in 0..<number.count {
      switch(number[i]) {
      case "[":
          depth += 1
          if depth < 5 {
              out.append("[")
          }
      case "]":
          depth -= 1
          out.append("]")
      case ",":
          out.append(",")
      default:
          if depth < 5 {
              out.append(number[i])
          } else {
              // assume number is correct, i.e. it has been exploded
              // before further nesting can happen
              // i: number, i+1: comma, i+2: number, i+3: ]
              let left = Int(number[i])!
              let right = Int(number[i+2])!
              for j in (0..<out.count).reversed() {
                  if let n = Int(out[j]) {
                      out[j] = String(n + left)
                      break
                  }
              }
              var rest = number[(i+4)...].map { $0 }
              for j in 0..<rest.count {
                  if let n = Int(rest[j]) {
                      rest[j] = String(n + right)
                      break
                  }
              }
              out = out + ["0"] + rest
              return FlatSnail(number: out)
          }
      }
  }
  return nil
}
That function returns nil if the number didn't change. The exploded pair is replaced by a 0, and the numbers of the pair are added to the closest number to the left (if any), and to the right (if any), respectively. This can be done by using indices (or "moving the cursor").

For completeness, here's the split function as well, which also returns an optional (nil if the number does not change):

func split() -> FlatSnail? {
    var out: [String] = []
    for i in 0..<number.count {
        switch(number[i]) {
        case "[":
            out.append("[")
        case "]":
            out.append("]")
        case ",":
            out.append(",")
        default:
            let n = Int(number[i])!
            if n >= 10 {
                let r = n % 2
                let left = String(n/2)
                let right = String((n+r)/2)
                let rest = number[(i+1)...].map { $0 }
                out = out + ["[", left, ",", right, "]"] + rest
                return FlatSnail(number: out)
            }
            out.append(number[i])
        }
    }
    return nil
}
And with this, it's now quite simple to complete the last part of the AoC exercise, which it's to compute the magnitude of these snailfish numbers. But I'll leave that to you as an exercise.

Conclusion

Swift enumerations are very powerful and can be used to represent even recursive structures like binary trees in a very compact manner. However, you won't have references to the elements of the tree, so you will need to recreate the tree every time you operate on it.

For the given problem, it turns out that it was easier to search for numbers linearly. We could flatten the tree (i.e. print the description), alter it, and parse it as a tree again. But we can simply work with a flattened tree all the time. Think on the operations you need to do on the data, and don't just simply buy the fanciest data structure in your language of choice. 🙈

P.S. Because of my bad choice, I couldn't solve the problem on the day, the 18th of December. I came back to it on Christmas day, and I redid it from scratch in an hour or so. In the Leaderboards, the fastest person has solved it in 17 minutes. 😳 I think I need at least 15 minutes to read the problem properly. 😅


◀️ Older | Newer ▶️

⏪ Previous year | Next year ⏩