Skip to content

Instantly share code, notes, and snippets.

@natecook1000 natecook1000/SortedCollection.swift Secret
Last active Jan 23, 2017

Embed
What would you like to do?
// MARK: Sorted Protocols
/// A collection with elements that are always in order according to its
/// `areInIncreasingOrder` property.
public protocol SortedCollection : Collection {
// FIXME: unsupported generic features
// associatedtype SubSequence : SortedCollection
var areInIncreasingOrder: (Iterator.Element, Iterator.Element) -> Bool { get }
/// Returns the insertion point for the given element.
///
/// The insertion point will be immediately after any elements that compare
/// as equal to `element` and before any elements that compare as greater
/// than `element`.
func stableInsertionPoint(for element: Iterator.Element) -> Index
}
/// A mutable, always-sorted collection.
public protocol MutableSortedCollection : SortedCollection {
/// Inserts the given element and returns its new position.
@discardableResult
mutating func insert(_ element: Iterator.Element) -> Index
/// Customization point for faster insertion.
///
/// The index passed as `i` must be the correct, stable insertion point for
/// `element`.
mutating func _insert(_ element: Iterator.Element, atSortedIndex i: Index)
/// Inserts the contents of `newElements` into the collection.
///
/// `newElements` does not need to be sorted before calling
/// `insert(contentsOf:)`.
mutating func insert<S: Sequence>(contentsOf newElements: S)
where S.Iterator.Element == Iterator.Element
/// Inserts the contents of the previously-sorted `newElements` sequence.
///
/// `newElements` must already be sorted using the same
/// `areInIncreasingOrder` predicate as this collection.
mutating func insert<S: Sequence>(sortedContentsOf newElements: S)
where S.Iterator.Element == Iterator.Element
/// Removes the element at `i`.
@discardableResult
mutating func remove(at i: Index) -> Iterator.Element
/// Removes the elements in the given subrange.
mutating func removeSubrange(_ bounds: Range<Index>)
/// Removes all elements from the collection.
mutating func removeAll(keepingCapacity: Bool)
}
// MARK: Default Implementations
extension SortedCollection {
public func stableInsertionPoint(for element: Iterator.Element) -> Index {
return _partitionPoint(where: { areInIncreasingOrder(element, $0) })
}
// Workaround for missing generic features.
public func _stableInsertionPoint(for element: Iterator.Element, startingAt i: Index) -> Index {
return _partitionPoint(where: { areInIncreasingOrder(element, $0) }, startingAt: i)
}
// Customization point for efficient `index(of:)`
public func _customIndexOfEquatableElement(_ element: Iterator.Element) -> Index?? {
let i = _partitionPoint(where: { !areInIncreasingOrder($0, element) })
guard i != endIndex
else { return .some(nil) }
return !areInIncreasingOrder(element, self[i])
? .some(i)
: .some(nil)
}
// Customization point for efficient `contains()`
public func _customContainsEquatableElement(_ element: Iterator.Element) -> Bool? {
return _customIndexOfEquatableElement(element)! != nil
}
}
extension MutableSortedCollection {
public mutating func insert<S: Sequence>(contentsOf newElements: S)
where S.Iterator.Element == Iterator.Element
{
for element in newElements {
insert(element)
}
}
public mutating func insert<S: Sequence>(sortedContentsOf newElements: S)
where S.Iterator.Element == Iterator.Element
{
var i = startIndex
for element in newElements {
i = _stableInsertionPoint(for: element, startingAt: i)
_insert(element, atSortedIndex: i)
}
}
@discardableResult
public mutating func remove(at i: Index) -> Iterator.Element {
let element = self[i]
removeSubrange(i..<index(after: i))
return element
}
public mutating func removeAll(keepingCapacity: Bool) {
removeSubrange(startIndex..<endIndex)
}
public mutating func _insert(_ element: Iterator.Element, atSortedIndex i: Index) {
insert(element)
}
}
// MARK: - Sorted Types
/// An always-sorted random-access array.
public struct SortedArray<Element>: MutableSortedCollection, RandomAccessCollection {
public let areInIncreasingOrder: (Element, Element) -> Bool
private var _elements: [Element]
// Collection requirements
public var startIndex: Int { return 0 }
public var endIndex: Int { return _elements.count }
public subscript(i: Int) -> Element {
return _elements[i]
}
public func index(after i: Int) -> Int {
return i + 1
}
public func index(before i: Int) -> Int {
return i - 1
}
public func index(_ i: Int, offsetBy n: Int) -> Int {
return i + n
}
public func distance(from start: Int, to end: Int) -> Int {
return end - start
}
/// Creates a new `SortedArray` by sorting the given sequence using the given
/// predicate.
public init<S: Sequence>(sorting elements: S,
areInIncreasingOrder: @escaping (Element, Element) -> Bool)
where S.Iterator.Element == Element
{
self._elements = elements.sorted(by: areInIncreasingOrder)
self.areInIncreasingOrder = areInIncreasingOrder
}
/// Creates a new `SortedArray` using the given sequence and predicate.
///
/// `elements` must already be in order before calling
/// `init(sorted:areInIncreasingOrder:)`.
public init<S: Sequence>(sorted elements: S,
areInIncreasingOrder: @escaping (Element, Element) -> Bool)
where S.Iterator.Element == Element,
S.SubSequence : Sequence,
S.SubSequence.Iterator.Element == Element
{
precondition(elements.isSorted(by: areInIncreasingOrder),
"unsorted elements passed to init(sorted:...)")
self._elements = Array(elements)
self.areInIncreasingOrder = areInIncreasingOrder
}
public mutating func _insert(_ element: Element, atSortedIndex i: Int) {
precondition(i == 0 || !areInIncreasingOrder(element, _elements[i - 1]),
"insert(_:atSortedIndex:) with incorrect index")
precondition(i >= _elements.count - 1 || areInIncreasingOrder(element, _elements[i + 1]),
"insert(_:atSortedIndex:) with incorrect index")
_elements.insert(element, at: i)
}
@discardableResult
public mutating func insert(_ element: Element) -> Int {
let i = stableInsertionPoint(for: element)
_elements.insert(element, at: i)
return i
}
public mutating func removeSubrange(_ bounds: Range<Int>) {
_elements.removeSubrange(bounds)
}
}
extension SortedArray where Element : Comparable {
/// Creates a new `SortedArray` by sorting the given sequence.
///
/// The new array uses the element type's default `<` ordering.
public init<S: Sequence>(sorting elements: S)
where S.Iterator.Element == Element
{
self = .init(sorting: elements, areInIncreasingOrder: <)
}
/// Creates a new `SortedArray` using the given sequence.
///
/// The new array uses the element type's default `<` ordering.
/// `elements` must already be in ascending order before calling
/// `init(sorted:)`.
public init<S: Sequence>(sorted elements: S)
where S.Iterator.Element == Element,
S.SubSequence : Sequence,
S.SubSequence.Iterator.Element == Element
{
self = .init(sorted: elements, areInIncreasingOrder: <)
}
}
/// A wrapper around an already-sorted collection.
public struct SortedView<Base: Collection> : SortedCollection
where Base.Iterator.Element == Base.SubSequence.Iterator.Element
{
typealias Element = Base.Iterator.Element
public typealias Index = Base.Index
public let areInIncreasingOrder: (Element, Element) -> Bool
let base: Base
// Collection requirements
public var startIndex: Index { return base.startIndex }
public var endIndex: Index { return base.endIndex }
public func index(after i: Index) -> Index {
return base.index(after: i)
}
public subscript(i: Index) -> Element {
return base[i]
}
/// Creates a new `SortedView` using the given sequence and predicate.
///
/// `elements` must already be in order before calling
/// `init(sorted:areInIncreasingOrder:)`.
public init(sorted elements: Base,
areInIncreasingOrder: @escaping (Element, Element) -> Bool)
{
precondition(elements.isSorted(by: areInIncreasingOrder),
"unsorted elements passed to init(sorted:...)")
self.base = elements
self.areInIncreasingOrder = areInIncreasingOrder
}
}
extension SortedView where Base.Iterator.Element : Comparable {
/// Creates a new `SortedView` using the given sequence.
///
/// The new wrapper uses the element type's default `<` ordering.
/// `elements` must already be in ascending order before calling
/// `init(sorted:)`.
public init(sorted elements: Base) {
self = .init(sorted: elements, areInIncreasingOrder: <)
}
}
// MARK: - Sequence and Collection Extensions
extension Sequence
where SubSequence : Sequence,
SubSequence.Iterator.Element == Iterator.Element
{
/// Returns `true` if the sequence is in order according to the given closure.
public func isSorted(
by areInIncreasingOrder: (Iterator.Element, Iterator.Element) -> Bool) -> Bool
{
for (a, b) in zip(self, self.dropFirst()) {
if areInIncreasingOrder(b, a) {
return false
}
}
return true
}
}
extension Collection {
/// Returns a partitioning index `p`, where `predicate` returns `false` for
/// all the elements up to, but not including `p`, and `true` for all
/// elements from `p` onward.
///
/// The collection must already be partitioned according to `predicate`.
internal func _partitionPoint(
where predicate: (Iterator.Element) throws -> Bool
) rethrows -> Index
{
return try _partitionPoint(where: predicate, startingAt: startIndex)
}
// Workaround for missing generic features.
internal func _partitionPoint(
where predicate: (Iterator.Element) throws -> Bool,
startingAt i: Index
) rethrows -> Index
{
var n = distance(from: i, to: endIndex)
var r = i
while n > 0 {
let half = n / 2
let mid = index(r, offsetBy: half)
if try predicate(self[mid]) {
n = half
}
else {
r = index(after: mid)
n -= half + 1
}
}
return r
}
}
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
You can’t perform that action at this time.