CountedSet.swift

/// An unordered collection of unique elements that
/// may appear more than once in the collection.
///
/// Each distinct value inserted into an `CountedSet`
/// has an associated counter. `CountedSet` keeps
/// track of the number of times each value is inserted.
/// That value must be removed the same number
/// of times.
///
/// Like `Set`, each `CountedSet` contains one entry for
/// each value, even when a value has been added multiple
/// times. The `count` method for both `Set` and `CountedSet`
/// returns the total number of values stored in each set.
/// A `CountedSet` does _not_ return the total number of times
/// each value is represented in the set.
///
/// New counted sets can be constructed from arrays.
/// Assign an array literal to a variable or constant
/// of the `CountedSet` type.
///
///    let quiver: CountedSet<Arrow> = [Arrow.iron, .elven, .elven]
///    if let count = quiver[.iron] {
///        print("There are \(count) iron arrows in your quiver")
///    }
///
/// Set Operations
/// ==============
///
/// Counted sets provide a suite of mathematical set operations.
/// For example, you can
/// efficiently test for membership of an element or check
/// intersection with another counted set:
///
/// - Use the `contains(_:)` method to test whether a set contains a specific
///   element.
/// - Use the "equal to" operator (`==`) to test whether two sets contain the
///   same elements.
/// - Use the `isSubset(of:)` method to test whether a set contains all the
///   elements of another set.
/// - Use the `isSuperset(of:)` method to test whether all elements of a set
///   are contained in another set.
/// - Use the `isDisjoint(with:)` method to test whether a set has any elements
///   in common with another set.
///
/// You can also combine, exclude, or subtract the elements of two counted sets:
///
/// - Use the `union(_:)` method to create a new set with the elements of a set
///   and another set or sequence.
/// - Use the `intersection(_:)` method to create a new set with only the
///   elements common to a set and another set or sequence.
/// - Use the `subtracting(_:)` method to create a new set with the elements of
///   a set that are not also in another set or sequence.
///
/// You can modify a counted set in place by using these methods' mutating
/// counterparts: `formUnion(_:)`, `formIntersection(_:)`,
/// and `subtract(_:)`.
///
public struct CountedSet<Element: Hashable> {
    /// Inserts the given element to the set, with
    /// a count of 1 if the element is not yet present
    /// or adding 1 to the count if it is.
    @discardableResult
    public mutating func insert(_ member: Element) -> Int {
        _counts[member, default: 0] += 1
        return _counts[member]!
    }
    
    /// Removes one instance of the given element from
    /// the set. If the count goes to zero, the element is
    /// removed from the set. If the element is not present,
    /// the request is ignored.
    @discardableResult
    public mutating func remove(_ member: Element) -> Int {
        let memberCount = _counts[member, default: 0]
        guard memberCount > 0 else { return 0 }
        guard memberCount > 1 else { _counts.removeValue(forKey: member); return 0 }
        let newCount = _counts[member]! - 1
        _counts[member] = newCount
        return newCount
    }
    
    /// Access an element count, returning 0 for any
    /// element that does not appear in the counted set
    public subscript(_ member: Element) -> Int {
        return _counts[member, default: 0]
    }
    
    /// Return the number of unique elements in the counted set
    public var count: Int { return _counts.keys.count }
    
    /// A Boolean value that indicates whether the counted set is empty.
    public var isEmpty: Bool {
        return count == 0
    }
    
    /// Return an element of the set
    public var anyElement: Element? {
        return _counts.keys.first
    }
    
    internal var _counts: [Element: Int]
}

extension CountedSet: Equatable {
    public static func ==<T>(lhs: CountedSet<T>, rhs: CountedSet<T>) -> Bool {
        return lhs._counts == rhs._counts
    }
}

extension CountedSet: Sequence {
    /// Returns an iterator over the members and counts
    /// of the counted set.
    public func makeIterator() -> DictionaryIterator<Element, Int> {
        return _counts.makeIterator()
    }

    /// Returns a Boolean value indicating whether
    /// the counted set contains a member
    public func contains(_ member: Element) -> Bool {
        return self[member] != 0
    }
    
    /// Check for both subset and equality relationship between
    /// a set and some sequence (which may itself be a `Set`).
    ///
    /// (isSubset: lhs ⊂ rhs, isEqual: lhs ⊂ rhs and |lhs| = |rhs|)
    internal func _compareSets<Element>(_ lhs: CountedSet<Element>, _ rhs: CountedSet<Element>)
        -> (isSubset: Bool, isEqual: Bool) {
            for member in lhs {
                // Would be better if it used key/count and not key/value
                // but that would require a custom iterator
                if rhs[member.key] != member.value {
                    return (false, false)
                }
            }
            return (true, lhs.count == rhs.count)
    }

    /// Returns a Boolean value that indicates whether this set is a subset of
    /// the given set.
    ///
    /// Set *A* is a subset of another set *B* if every member of *A* is also a
    /// member of *B*.
    ///
    /// - Parameter other: Another set.
    /// - Returns: `true` if the set is a subset of `other`; otherwise, `false`.
    public func isSubset(of other: CountedSet<Element>) -> Bool {
        let (isSubset, isEqual) = _compareSets(self, other)
        return isSubset || isEqual
    }
    
    /// Returns a Boolean value that indicates whether the set is a superset of
    /// the given sequence.
    ///
    /// Set *A* is a superset of another set *B* if every member of *B* is also a
    /// member of *A*.
    ///
    /// - Parameter possibleSubset: another counted set
    /// - Returns: `true` if the set is a superset of `possibleSubset`;
    ///   otherwise, `false`.
    public func isSuperset(of other: CountedSet<Element>) -> Bool {
            return other.isSubset(of: self)
    }

    /// Returns a Boolean value that indicates whether this set has no members in
    /// common with the given set.
    ///
    /// - Parameter other: Another counted set.
    /// - Returns: `true` if the set has no elements in common with `other`;
    ///   otherwise, `false`.
    public func isDisjoint(with other: CountedSet<Element>) -> Bool {
        for (key, _) in self {
            if other[key] > 0 { return false }
        }
        return true
    }
    
    /// Inserts the elements of the given sequence into the set.
    ///
    /// If the set already contains one or more elements that are also in
    /// `other`, the existing members are kept. If `other` contains multiple
    /// instances of equivalent elements, only the first instance is kept.
    ///
    ///     var attendees: Set = ["Alicia", "Bethany", "Diana"]
    ///     let visitors = ["Diana", "Marcia", "Nathaniel"]
    ///     attendees.formUnion(visitors)
    ///     print(attendees)
    ///     // Prints "["Diana", "Nathaniel", "Bethany", "Alicia", "Marcia"]"
    ///
    /// - Parameter other: A sequence of elements. `other` must be finite.
    public mutating func formUnion(_ other: CountedSet<Element>) {
        for (item, count) in other {
            _counts[item] = count + self[item]
        }
    }
    
    /// Returns a new set with the elements of both this set and the given
    /// sequence.
    ///
    /// - Parameter other: Another counted set.
    /// - Returns: A new set with the unique elements of this set and `other`.
    public func union(_ other: CountedSet<Element>) -> CountedSet<Element> {
            var newSet = self
            newSet.formUnion(other)
            return newSet
    }
    
    /// Removes the elements of the set that aren't also in another set.
    ///
    /// - Parameter other: another counted set.
    public mutating func formIntersection(_ other: CountedSet<Element>) {
        for (key, _) in self {
            if other[key] == 0 {
                _counts.removeValue(forKey: key)
            }
        }
    }
    
    /// Returns a new set with the elements that are common to both this set and
    /// the other set.
    ///
    /// - Parameter other: Another counted set.
    /// - Returns: A new counted set.
    public func intersection(_ other: CountedSet<Element>) -> CountedSet<Element> {
        var newSet = self
        newSet.formIntersection(other)
        return newSet
    }
    
    /// Removes the elements of the set that are found in another set.
    ///
    /// - Parameter other: another counted set.
    public mutating func subtract(_ other: CountedSet<Element>) {
        for (key, _) in other {
            _counts.removeValue(forKey: key)
        }
    }
    
    /// Returns a new set with the elements that do not appear in
    /// another set.
    ///
    /// - Parameter other: Another counted set.
    /// - Returns: A new counted set.
    public func subtracting(_ other: CountedSet<Element>) -> CountedSet<Element> {
        var newSet = self
        newSet.subtract(other)
        return newSet
    }

}

extension CountedSet: ExpressibleByArrayLiteral {
    //
    // `ExpressibleByArrayLiteral` conformance
    //
    
    /// Creates a counted set containing the elements of the given array literal.
    ///
    /// Do not call this initializer directly. It is used by the compiler when
    /// you use an array literal. Instead, create a new counted set using an array
    /// literal as its value by enclosing a comma-separated list of values in
    /// square brackets. You can use an array literal anywhere a set is expected
    /// by the type context.
    public init(arrayLiteral elements: Element...) {
        _counts = [:]
        for element in elements {
            _counts[element] = _counts[element, default: 0] + 1
        }
    }
}

enum Arrow { case iron, wooden, elven, dwarvish, magic, silver }
var aCountedSet = CountedSet<Arrow>()
aCountedSet[.iron] // 0

var myCountedSet: CountedSet<Arrow> = [.iron, .magic, .iron, .silver, .iron, .iron]
myCountedSet[.iron] // 4

for (key, count) in myCountedSet {
    print("\(key): \(count)")
}
// silver: 1
// iron: 4
// magic: 1

myCountedSet.remove(.iron) // 3
myCountedSet.remove(.dwarvish) // 0
myCountedSet.remove(.magic) // 0

// magic should not appear here
for (key, count) in myCountedSet {
    print("\(key): \(count)")
}
// silver: 1
// iron: 3

myCountedSet.insert(.iron)
myCountedSet.contains(Arrow.iron)