Search code examples
swiftgenerics

How can I implement a CountedSet (NSCountedSet) in Swift?


Create a generic CountedSet struct that is constrained to Hashable elements. A counted set is an unordered collection of unique elements that may appear more than once in the collection. Use a private dictionary as your backing storage for set members and their counts.

struct CountedSet<Element> {
    private(set) var elements: [Element]
    
    mutating func insert(_ element: Element) {
        elements.append(element)
    }
    
    mutating func remove() -> Element? {
        guard elements.isEmpty == false else { return nil}
        return elements.removeFirst()
        
    }
    
    subscript(_ member: Element) -> Int {
        return 0
    }
}

I don't understand what the real objective is here. The instructions are very confusing at least to me.


Solution

  • 1) Make your generic struct element conform to Hashable, this is necessary because the dictionary keys are required to conform to Hashable.

    struct CountedSet<Element: Hashable>
    

    2) The backing storage you have used is an ordered array, not a dictionary and you need to initialize it with an empty one.

    private(set) var elements: [Element: Int] = [:]
    

    3) Your subscript method you need to return the count for the counted set member or zero if it is nil.

    return elements[member] ?? 0
    

    4) Your Insert and Remove methods need to first check the count of a member in the backing dictionary before adding or removing an element from it.

    So your CountedSet should look like this:

    struct CountedSet<Element: Hashable> {
        private(set) var elements: [Element: Int] = [:]
        mutating func insert(_ member: Element) {
            elements[member, default: 0] += 1
        }
        mutating func remove(_ member: Element) -> Element? {
            guard var count = elements[member], count > 0 else { return nil }
            count -= 1
            elements[member] = count == 0 ? nil : count
            return member
        }
        subscript(_ member: Element) -> Int {
            elements[member] ?? 0
        }
    }
    

    var countedSet = CountedSet<Int>()
    countedSet.insert(3)
    countedSet.insert(3)
    countedSet.insert(4)
    countedSet.elements  // [4: 1, 3: 2]
    countedSet.remove(4)
    countedSet.elements  // [3: 2]
    countedSet.remove(4) // nil
    

    Expanding on that you can also make your CountedSet conform to ExpressibleByArrayLiteral to allow you to initialize your CountedSet with an array and CustomStringConvertible to allow you to print its elements:

    extension CountedSet: ExpressibleByArrayLiteral, CustomStringConvertible {
    
        typealias ArrayLiteralElement = Element
        init<S: Sequence>(_ sequence: S) where S.Element == Element {
            self.elements = sequence.reduce(into: [:]) { $0[$1, default: 0] += 1 }
        }
        init(arrayLiteral elements: Element...)  { self.init(elements) }
        var description: String { .init(describing: elements) }
    }
    

    var countedSet: CountedSet = [1,2,2,3,3,3,4,4,5,5,5]
    print(countedSet) // "[5: 3, 2: 2, 3: 3, 4: 2, 1: 1]\n"