A Swift solver for the numbers game in Countdown

main.swift

// This is a solver for the Numbers game in the long-running British TV show "Countdown".
// It tries to find all possible solutions, discarding uninteresting variations etc.
//
// https://en.wikipedia.org/wiki/Countdown_(game_show)#Numbers_round

import ArgumentParser

struct Expr: Comparable, Hashable, CustomStringConvertible {
  enum Kind: Int {
  case value = 0
  case add = 1
  case mul = 2
  }

  var kind: Kind
  var v: Int
  var a: [Expr]
  var b: [Expr]

  init?(value: Int) {
    guard value > 0 else { return nil }
    self.kind = .value
    self.v = value
    self.a = []
    self.b = []
  }

  init?(add x: Expr, _ y: Expr) {
    self.kind = .add
    self.v = x.v + y.v
    switch (x.kind, y.kind) {
    case (.add, .add):
      self.a = (x.a + y.a).sorted()
      self.b = (x.b + y.b).sorted()
    case (.add, _):
      self.a = (x.a + [y]).sorted()
      self.b = x.b
    case (_, .add):
      self.a = ([x] + y.a).sorted()
      self.b = y.b
    default:
      self.a = [x, y].sorted()
      self.b = []
    }
  }

  init?(sub x: Expr, _ y: Expr) {
    self.kind = .add
    self.v = x.v - y.v
    guard v > 0 else { return nil }
    switch (x.kind, y.kind) {
    case (.add, .add):
      self.a = (x.a + y.b).sorted()
      self.b = (y.a + x.b).sorted()
    case (.add, _):
      self.a = x.a
      self.b = (x.b + [y]).sorted()
    case (_, .add):
      self.a = ([x] + y.b).sorted()
      self.b = y.a
    default:
      self.a = [x]
      self.b = [y]
    }
  }

  init?(mul x: Expr, _ y: Expr) {
    self.kind = .mul
    self.v = x.v * y.v
    switch (x.kind, y.kind) {
    case (.mul, .mul):
      self.a = (x.a + y.a).sorted()
      self.b = (x.b + y.b).sorted()
    case (.mul, _):
      self.a = (x.a + [y]).sorted()
      self.b = x.b
    case (_, .mul):
      self.a = ([x] + y.a).sorted()
      self.b = y.b
    default:
      self.a = [x, y].sorted()
      self.b = []
    }
  }

  init?(div x: Expr, _ y: Expr) {
    self.kind = .mul
    if y.v == 0 || x.v % y.v != 0 { return nil }
    self.v = x.v / y.v
    guard v > 0 else { return nil }
    switch (x.kind, y.kind) {
    case (.mul, .mul):
      self.a = (x.a + y.b).sorted()
      self.b = (y.a + x.b).sorted()
    case (.mul, _):
      self.a = x.a
      self.b = (x.b + [y]).sorted()
    case (_, .mul):
      self.a = ([x] + y.b).sorted()
      self.b = y.a
    default:
      self.a = [x]
      self.b = [y]
    }
  }

  static func == (left: Expr, right: Expr) -> Bool {
    guard left.kind == right.kind else { return false }
    guard left.v == right.v else { return false }
    return left.a == right.a && left.b == right.b
  }

  static func < (left: Expr, right: Expr) -> Bool {
    if left.kind.rawValue < right.kind.rawValue { return true }
    guard left.kind == right.kind else { return false }

    if left.kind == .value {
      return left.v < right.v
    }

    if left.a.lexicographicallyPrecedes(right.a) {
      return true
    }
    guard left.a == right.a else { return false }
    return left.b.lexicographicallyPrecedes(right.b)
  }

  var description: String {
    switch kind {
    case .value:
      return "\(v)"
    case .add:
      let s1 = a.map { "\($0)" }.joined(separator: " + ")
      let s2 = b.map { "\($0)" }.joined(separator: " - ")
      if b.isEmpty {
        return "(\(s1))"
      }
      return "(\(s1) - \(s2))"
    case .mul:
      let s1 = a.map { "\($0)" }.joined(separator: " * ")
      let s2 = b.map { "\($0)" }.joined(separator: " / ")
      if b.isEmpty {
        return "\(s1)"
      }
      return "\(s1) / \(s2)"
    }
  }
}

enum Op: String, CaseIterable {
  case add
  case sub
  case mul
  case div

  func apply(_ a: Expr, _ b: Expr) -> Expr? {
    switch self {
    case .add:
      return Expr(add: a, b)
    case .sub:
      return Expr(sub: a, b)
    case .mul:
      return Expr(mul: a, b)
    case .div:
      return Expr(div: a, b)
    }
  }
}

func countdown(input: [Int], target: Int, match: (Expr) -> Void) {
  var input: [Expr?] = input.map { Expr(value: $0) }
  for i in input.indices {
    if input[i]?.v == target {
      match(Expr(value: target)!)
      input[i] = nil
    }
  }
  var count = input.count
  _countdown(input: &input, count: &count, target: target, match: match)
}

func _countdown(
  input: inout [Expr?],
  count: inout Int,
  target: Int,
  match: (Expr) -> Void
) {
  if count <= 1 { return }
  for i in input.indices {
    guard let a = input[i] else { continue }
    input[i] = nil
    count -= 1
    defer { input[i] = a; count += 1 }
    for j in input.indices {
      guard let b = input[j] else { continue }
      for op in Op.allCases {
        if let c = op.apply(a, b) {
          input[j] = nil
          if c.v == target,
             input.lazy.compactMap({ $0?.kind != .value ? $0 : nil }).isEmpty
          {
            match(c)
          }
          input[j] = c
          _countdown(input: &input, count: &count, target: target, match: match)
          input[j] = b
        }
      }
    }
  }
}


struct CountdownSolver: ParsableCommand {
  @Option
  var target: Int

  @Option(parsing: .remaining)
  var input: [Int]

  func run() {
    var matches: Set<Expr> = []
    countdown(input: input, target: target) { trace in
      matches.insert(trace)
    }

    for match in matches.sorted() {
      print("\(target) = \(match)")
    }
    print("Number of distinct solutions: \(matches.count)")
  }
}

CountdownSolver.main()

Package.swift

// swift-tools-version:5.3
import PackageDescription
let package = Package(
  name: "countdown-solver",
  products: [
    .executable(name: "countdown-solver", targets: ["countdown-solver"]),
  ],
  dependencies: [
    .package(url: "https://github.com/apple/swift-argument-parser", from: "0.3.0"),
  ],
  targets: [
    .target(
      name: "countdown-solver",
      dependencies: [
        .product(name: "ArgumentParser", package: "swift-argument-parser"),
      ],
      path: "."),
  ]
)