Hex Swift

Hex.swift

// from: https://github.com/raywenderlich/swift-algorithm-club
public struct Heap<T> {
    
    /** The array that stores the heap's nodes. */
    var nodes = [T]()
    
    /**
     * Determines how to compare two nodes in the heap.
     * Use '>' for a max-heap or '<' for a min-heap,
     * or provide a comparing method if the heap is made
     * of custom elements, for example tuples.
     */
    private var orderCriteria: (T, T) -> Bool
    
    /**
     * Creates an empty heap.
     * The sort function determines whether this is a min-heap or max-heap.
     * For comparable data types, > makes a max-heap, < makes a min-heap.
     */
    public init(sort: @escaping (T, T) -> Bool) {
        self.orderCriteria = sort
    }
    
    /**
     * Creates a heap from an array. The order of the array does not matter;
     * the elements are inserted into the heap in the order determined by the
     * sort function. For comparable data types, '>' makes a max-heap,
     * '<' makes a min-heap.
     */
    public init(array: [T], sort: @escaping (T, T) -> Bool) {
        self.orderCriteria = sort
        configureHeap(from: array)
    }
    
    /**
     * Configures the max-heap or min-heap from an array, in a bottom-up manner.
     * Performance: This runs pretty much in O(n).
     */
    private mutating func configureHeap(from array: [T]) {
        nodes = array
        for i in stride(from: (nodes.count/2-1), through: 0, by: -1) {
            shiftDown(i)
        }
    }
    
    public var isEmpty: Bool {
        return nodes.isEmpty
    }
    
    public var count: Int {
        return nodes.count
    }
    
    /**
     * Returns the index of the parent of the element at index i.
     * The element at index 0 is the root of the tree and has no parent.
     */
    @inline(__always) internal func parentIndex(ofIndex i: Int) -> Int {
        return (i - 1) / 2
    }
    
    /**
     * Returns the index of the left child of the element at index i.
     * Note that this index can be greater than the heap size, in which case
     * there is no left child.
     */
    @inline(__always) internal func leftChildIndex(ofIndex i: Int) -> Int {
        return 2*i + 1
    }
    
    /**
     * Returns the index of the right child of the element at index i.
     * Note that this index can be greater than the heap size, in which case
     * there is no right child.
     */
    @inline(__always) internal func rightChildIndex(ofIndex i: Int) -> Int {
        return 2*i + 2
    }
    
    /**
     * Returns the maximum value in the heap (for a max-heap) or the minimum
     * value (for a min-heap).
     */
    public func peek() -> T? {
        return nodes.first
    }
    
    /**
     * Adds a new value to the heap. This reorders the heap so that the max-heap
     * or min-heap property still holds. Performance: O(log n).
     */
    public mutating func insert(_ value: T) {
        nodes.append(value)
        shiftUp(nodes.count - 1)
    }
    
    /**
     * Adds a sequence of values to the heap. This reorders the heap so that
     * the max-heap or min-heap property still holds. Performance: O(log n).
     */
    public mutating func insert<S: Sequence>(_ sequence: S) where S.Iterator.Element == T {
        for value in sequence {
            insert(value)
        }
    }
    
    /**
     * Allows you to change an element. This reorders the heap so that
     * the max-heap or min-heap property still holds.
     */
    public mutating func replace(index i: Int, value: T) {
        guard i < nodes.count else { return }
        
        remove(at: i)
        insert(value)
    }
    
    /**
     * Removes the root node from the heap. For a max-heap, this is the maximum
     * value; for a min-heap it is the minimum value. Performance: O(log n).
     */
    @discardableResult public mutating func remove() -> T? {
        guard !nodes.isEmpty else { return nil }
        
        if nodes.count == 1 {
            return nodes.removeLast()
        } else {
            // Use the last node to replace the first one, then fix the heap by
            // shifting this new first node into its proper position.
            let value = nodes[0]
            nodes[0] = nodes.removeLast()
            shiftDown(0)
            return value
        }
    }
    
    /**
     * Removes an arbitrary node from the heap. Performance: O(log n).
     * Note that you need to know the node's index.
     */
    @discardableResult public mutating func remove(at index: Int) -> T? {
        guard index < nodes.count else { return nil }
        
        let size = nodes.count - 1
        if index != size {
            nodes.swapAt(index, size)
            shiftDown(from: index, until: size)
            shiftUp(index)
        }
        return nodes.removeLast()
    }
    
    /**
     * Takes a child node and looks at its parents; if a parent is not larger
     * (max-heap) or not smaller (min-heap) than the child, we exchange them.
     */
    internal mutating func shiftUp(_ index: Int) {
        var childIndex = index
        let child = nodes[childIndex]
        var parentIndex = self.parentIndex(ofIndex: childIndex)
        
        while childIndex > 0 && orderCriteria(child, nodes[parentIndex]) {
            nodes[childIndex] = nodes[parentIndex]
            childIndex = parentIndex
            parentIndex = self.parentIndex(ofIndex: childIndex)
        }
        
        nodes[childIndex] = child
    }
    
    /**
     * Looks at a parent node and makes sure it is still larger (max-heap) or
     * smaller (min-heap) than its childeren.
     */
    internal mutating func shiftDown(from index: Int, until endIndex: Int) {
        let leftChildIndex = self.leftChildIndex(ofIndex: index)
        let rightChildIndex = leftChildIndex + 1
        
        // Figure out which comes first if we order them by the sort function:
        // the parent, the left child, or the right child. If the parent comes
        // first, we're done. If not, that element is out-of-place and we make
        // it "float down" the tree until the heap property is restored.
        var first = index
        if leftChildIndex < endIndex && orderCriteria(nodes[leftChildIndex], nodes[first]) {
            first = leftChildIndex
        }
        if rightChildIndex < endIndex && orderCriteria(nodes[rightChildIndex], nodes[first]) {
            first = rightChildIndex
        }
        if first == index { return }
        
        nodes.swapAt(index, first)
        shiftDown(from: first, until: endIndex)
    }
    
    internal mutating func shiftDown(_ index: Int) {
        shiftDown(from: index, until: nodes.count)
    }
    
}

// MARK: - Searching
extension Heap where T: Equatable {
    
    /** Get the index of a node in the heap. Performance: O(n). */
    public func index(of node: T) -> Int? {
        return nodes.firstIndex(where: { $0 == node })
    }
    
    /** Removes the first occurrence of a node from the heap. Performance: O(n log n). */
    @discardableResult public mutating func remove(node: T) -> T? {
        if let index = index(of: node) {
            return remove(at: index)
        }
        return nil
    }
    
}

// from: https://github.com/raywenderlich/swift-algorithm-club
public struct PriorityQueue<T> {
    fileprivate var heap: Heap<T>
    
    /*
     To create a max-priority queue, supply a > sort function. For a min-priority
     queue, use <.
     */
    public init(sort: @escaping (T, T) -> Bool) {
        heap = Heap(sort: sort)
    }
    
    public var isEmpty: Bool {
        return heap.isEmpty
    }
    
    public var count: Int {
        return heap.count
    }
    
    public func peek() -> T? {
        return heap.peek()
    }
    
    public mutating func enqueue(_ element: T) {
        heap.insert(element)
    }
    
    public mutating func dequeue() -> T? {
        return heap.remove()
    }
    
    /*
     Allows you to change the priority of an element. In a max-priority queue,
     the new priority should be larger than the old one; in a min-priority queue
     it should be smaller.
     */
    public mutating func changePriority(index i: Int, value: T) {
        return heap.replace(index: i, value: value)
    }
}

extension PriorityQueue where T: Equatable {
    public func index(of element: T) -> Int? {
        return heap.index(of: element)
    }
}

//////////////////////////////////////////////////////////////////////
//////////// https://www.redblobgames.com/grids/hexagons/ ////////////
//////////////////////////////////////////////////////////////////////
import Foundation

public struct HexCoordinate {
    public var x: Int
    public var y: Int
    public var z: Int
    
    public init(x: Int, y: Int, z: Int) {
        assert(x + y + z == 0, "Invalid coordinate")
        self.x = x
        self.y = y
        self.z = z
    }
    
    public init(x: Int, y: Int) {
        self.init(x: x, y: y, z: -x-y)
    }
    
    public init(x: Int, z: Int) {
        self.init(x: x, y: -x-z, z: z)
    }
    
    public init(y: Int, z: Int) {
        self.init(x: -y-z, y: y, z: z)
    }
}

public extension HexCoordinate {
    
    func withOffset(of hex: HexCoordinate) -> HexCoordinate {
        return .init(x: self.x + hex.x, y: self.y + hex.y, z: self.z + hex.z)
    }
    
    func withOffset(x: Int, y: Int, z: Int) -> HexCoordinate {
        return .init(x: self.x + x, y: self.y + y, z: self.z + z)
    }
    
    func withOffset(x: Int, y: Int) -> HexCoordinate {
        return .init(x: self.x + x, y: self.y + y)
    }
    
    func withOffset(x: Int, z: Int) -> HexCoordinate {
        return .init(x: self.x + x, z: self.z + z)
    }
    
    func withOffset(y: Int, z: Int) -> HexCoordinate {
        return .init(y: self.y + y, z: self.z + z)
    }
}

extension HexCoordinate: Equatable {
    public static func ==(lhs: HexCoordinate, rhs: HexCoordinate) -> Bool {
        return lhs.x == rhs.x && lhs.z == rhs.z // y has to match if x and z match, no need to check
    }
}

extension HexCoordinate: Hashable {}

extension HexCoordinate: CustomStringConvertible {
    public var description: String {
        return "{x:\(x), y:\(y), z:\(z)}"
    }
}

extension HexCoordinate: CustomDebugStringConvertible {
    public var debugDescription: String {
        return "Hex(x:\(x), y:\(y), z:\(z))"
    }
}

public class Hex<Storage> {
    public let coordinate: HexCoordinate
    public var storage: Storage? = nil
    
    internal init(coordinate: HexCoordinate) {
        self.coordinate = coordinate
    }
}

extension Hex: CustomStringConvertible {
    public var description: String {
        return "Hex[\(coordinate)]"
    }
}

extension Hex: CustomDebugStringConvertible {
    public var debugDescription: String {
        return "Hex[\(coordinate)] storage: \(String(describing: storage))"
    }
}

public enum HexValue {
    case value(Int)
    case unavailable
}

extension HexValue {
    public static func +(lhs: HexValue, rhs: HexValue) -> HexValue {
        switch (lhs, rhs) {
        case (.unavailable, _), (_, .unavailable):
            return .unavailable
        case let (.value(lval), .value(rval)):
            return .value(lval + rval)
        }
    }
}

extension HexValue: Comparable {
    public static func < (lhs: HexValue, rhs: HexValue) -> Bool {
        switch (lhs, rhs) {
        case (.unavailable, _):
            return false
        case (_, .unavailable):
            return true
        case let (.value(lval), .value(rval)):
            return lval < rval
        }
    }
    
}

public typealias HexCost<Storage> = (Hex<Storage>) -> HexValue

public final class HexGrid<Storage> {
    private var storage: ContiguousArray<ContiguousArray<Hex<Storage>>>
    public let width: Int
    public let height: Int
    
    public init(width: Int, height: Int) {
        assert(width >= 0)
        assert(height >= 0)
        self.width = width
        self.height = height
        self.storage = .init()
        initializeFields(width: width, height: height)
    }
    
    private func initializeFields(width: Int, height: Int) {
        self.storage.reserveCapacity(width)
        (0..<width).forEach { c in
            var column: ContiguousArray<Hex<Storage>> = .init()
            column.reserveCapacity(height)
            (0..<width).forEach { r in
                column.append(.init(coordinate: HexCoordinate(x: c, z: r)))
            }
            self.storage.append(column)
        }
    }
    
    public func isValid(coordinate: HexCoordinate) -> Bool {
        return coordinate.x >= 0 && coordinate.x < width && coordinate.z >= 0 && coordinate.z < height
    }
}

// MARK: access
public extension HexGrid {
    subscript(coordinate: HexCoordinate) -> Hex<Storage> {
        assert(isValid(coordinate: coordinate), "Invalid coordinate")
        return self.storage[coordinate.x][coordinate.z]
    }
}

// MARK: neighborhood
public extension HexGrid {
    
    func neighborhood(of coordinate: HexCoordinate) -> [HexCoordinate] {
        let indices = [
            coordinate.withOffset(x: 1, y: -1, z: 0),
            coordinate.withOffset(x: 1, y: 0, z: -1),
            coordinate.withOffset(x: 0, y: 1, z: -1),
            coordinate.withOffset(x: -1, y: 1, z: 0),
            coordinate.withOffset(x: -1, y: 0, z: 1),
            coordinate.withOffset(x: 0, y: -1, z: 1),
        ]
        
        return indices.filter(self.isValid)
    }
    
    func isNeighbor(of coordinate: HexCoordinate) -> (HexCoordinate) -> Bool {
        return { other in
            self.neighborhood(of: coordinate).contains(other)
        }
    }
}

// MARK: diagonals
public extension HexGrid {

    func diagonals(of coordinate: HexCoordinate) -> [HexCoordinate] {
        let indices = [
            coordinate.withOffset(x: 2, y: -1, z: -1),
            coordinate.withOffset(x: 1, y: 1, z: -2),
            coordinate.withOffset(x: -1, y: 2, z: -1),
            coordinate.withOffset(x: -2, y: 1, z: 1),
            coordinate.withOffset(x: -1, y: -1, z: 2),
            coordinate.withOffset(x: 1, y: -2, z: 1),
        ]
        
        return indices.filter(self.isValid)
    }
    
    func isDiagonal(of coordinate: HexCoordinate) -> (HexCoordinate) -> Bool {
        return { other in
            self.diagonals(of: coordinate).contains(other)
        }
    }
}

// MARK: range
public extension HexGrid {
    func range(of size: Int) -> (HexCoordinate) -> [HexCoordinate] {
        assert(size >= 0)
        return { center in
            var indices: Array<HexCoordinate> = .init()
            (-size...size).forEach { idx in
                (-size...size).forEach { idy in
                    (-size...size).forEach { idz in
                        if idx + idy + idz == 0 {
                            indices.append(
                                HexCoordinate(x: idx + center.x, y: idy + center.y, z: idz + center.z)
                            )
                        } else { /* nothing */ }
                    }
                }
            }
            
            return indices.filter(self.isValid)
        }
    }
    
    func isInRange(of size: Int) -> (HexCoordinate) -> (HexCoordinate) -> Bool {
        assert(size >= 0)
        return { center in
            return { other in
                // it might be optimized to just check indicies
                return self.range(of: size)(center).contains(other)
            }
        }
    }
    
    func weightedRange(of size: Int, cost: @escaping HexCost<Storage>) -> (HexCoordinate) -> [HexCoordinate] {
        fatalError()
    }
    
    func isInWeightedRange(of size: Int, cost: @escaping HexCost<Storage>) -> (HexCoordinate) -> (HexCoordinate) -> Bool {
        assert(size >= 0)
        return { center in
            return { other in
                return self.weightedRange(of: size, cost: cost)(center).contains(other)
            }
        }
    }
}

// MARK: range intersection
public extension HexGrid {
    
    /*
     var results = []
     for each xmin ≤ x ≤ xmax:
     for each max(ymin, -x-zmax) ≤ y ≤ min(ymax, -x-zmin):
     var z = -x-y
     results.append(Cube(x, y, z))
     */
}

// MARK: line
public extension HexGrid {
    
    func line(from: HexCoordinate, to: HexCoordinate) -> [HexCoordinate] {
        assert(isValid(coordinate: from))
        assert(isValid(coordinate: to))
        let dist = Double(distance(from: from, to: to))
        let indices: Array<HexCoordinate>
            = (0...Int(dist)).map { (idx) -> HexCoordinate in
                let i = Double(idx)
                return roundToHex(
                    x:lerp(a: Double(from.x), b: Double(to.x), t: 1.0 / dist * i),
                    y:lerp(a: Double(from.y), b: Double(to.y), t: 1.0 / dist * i),
                    z:lerp(a: Double(from.z), b: Double(to.z), t: 1.0 / dist * i)
                )
        }
        return indices.filter(self.isValid)
    }
    
    private func lerp(a: Double, b: Double, t: Double) -> Double {
        return a + (b - a) * t
    }
    
    private func roundToHex(x: Double, y: Double, z: Double) -> HexCoordinate {
        var rx = x.rounded()
        var ry = y.rounded()
        var rz = z.rounded()
        
        let x_diff = abs(rx - x)
        let y_diff = abs(ry - y)
        let z_diff = abs(rz - z)
        
        if x_diff > y_diff && x_diff > z_diff {
            rx = -ry-rz
        } else if y_diff > z_diff {
            ry = -rx-rz
        } else {
            rz = -rx-ry
        }
        return HexCoordinate.init(x: Int(rx), y: Int(ry), z: Int(rz))
    }
}


// MARK: distance
public extension HexGrid {

    func distance(from: HexCoordinate, to: HexCoordinate) -> Int {
        assert(isValid(coordinate: from))
        assert(isValid(coordinate: to))
        return max(abs(from.x - to.x), abs(from.y - to.y), abs(from.z - to.z))
    }
    
    func weightedDistance(from: HexCoordinate, to: HexCoordinate, cost: @escaping HexCost<Storage>) -> HexValue {
        assert(isValid(coordinate: from))
        assert(isValid(coordinate: to))
        return .value(distance(from: from, to: to)) // TODO: change to valid implementation...
    }
}

// MARK: path
public extension HexGrid {
    /*
     function hex_reachable(start, movement):
     var visited = set() # set of hexes
     add start to visited
     var fringes = [] # array of arrays of hexes
     fringes.append([start])
     
     for each 1 < k ≤ movement:
     fringes.append([])
     for each hex in fringes[k-1]:
     for each 0 ≤ dir < 6:
     var neighbor = hex_neighbor(hex, dir)
     if neighbor not in visited and not blocked:
     add neighbor to visited
     fringes[k].append(neighbor)
     
     return visited
     */
    
    /*
     frontier = PriorityQueue()
     frontier.put(start, 0)
     came_from = {}
     cost_so_far = {}
     came_from[start] = None
     cost_so_far[start] = 0
     
     while not frontier.empty():
     current = frontier.get()
     
     if current == goal:
     break
     
     for next in graph.neighbors(current):
     new_cost = cost_so_far[current] + graph.cost(current, next)
     if next not in cost_so_far or new_cost < cost_so_far[next]:
     cost_so_far[next] = new_cost
     priority = new_cost + heuristic(goal, next)
     frontier.put(next, priority)
     came_from[next] = current
     */
    
    /*
     def heuristic(a, b):
     (x1, y1) = a
     (x2, y2) = b
     return abs(x1 - x2) + abs(y1 - y2)
     
     def a_star_search(graph, start, goal):
     frontier = PriorityQueue()
     frontier.put(start, 0)
     came_from = {}
     cost_so_far = {}
     came_from[start] = None
     cost_so_far[start] = 0
     
     while not frontier.empty():
     current = frontier.get()
     
     if current == goal:
     break
     
     for next in graph.neighbors(current):
     new_cost = cost_so_far[current] + graph.cost(current, next)
     if next not in cost_so_far or new_cost < cost_so_far[next]:
     cost_so_far[next] = new_cost
     priority = new_cost + heuristic(goal, next)
     frontier.put(next, priority)
     came_from[next] = current
     
     return came_from, cost_so_far
     
     def reconstruct_path(came_from, start, goal):
     current = goal
     path = []
     while current != start:
     path.append(current)
     current = came_from[current]
     path.append(start) # optional
     path.reverse() # optional
     return path
     */
    
    func path(from: HexCoordinate, to: HexCoordinate, cost: @escaping HexCost<Storage> = { _ in .value(1) }) -> (cost: HexValue, path: [HexCoordinate]) {
        assert(isValid(coordinate: from))
        assert(isValid(coordinate: to))
        var queue: PriorityQueue<(priority: HexValue, coordinate: HexCoordinate)>
            = .init { (lhs, rhs) -> Bool in return lhs.priority < rhs.priority }
        var from_coord: [HexCoordinate:HexCoordinate] = [:]
        var cost_sum: [HexCoordinate:HexValue] = [:]
        
        queue.enqueue((.value(0), from))
        cost_sum[from] = .value(0)
        
        while let current = queue.dequeue() {
            for next in neighborhood(of: current.coordinate) {
                let new_cost = (cost_sum[current.coordinate] ?? .unavailable) + cost(self[next])
                if new_cost < (cost_sum[next] ?? .unavailable) {
                    cost_sum[next] = new_cost
                    let priority = new_cost + weightedDistance(from: next, to: to, cost: cost)
                    queue.enqueue((priority, next))
                    from_coord[next] = current.coordinate
                } else { /* nothing */ }
            }
        }
        
        var path: [HexCoordinate] = []
        var current: HexCoordinate = to
        while let next = from_coord[current] {
            path.append(current)
            current = next
        }
        path.append(from)
        
        return (cost: cost_sum[to] ?? .unavailable, path: path.reversed()) // optional reverse
    }
}

extension HexCoordinate: CustomPlaygroundDisplayConvertible {
    public var playgroundDescription: Any {
        return self.debugDescription
    }
}
extension Hex: CustomPlaygroundDisplayConvertible {
    public var playgroundDescription: Any {
        return self.debugDescription
    }
}