|
2 | 2 | // |
3 | 3 | // This source file is part of the Swift Algorithms open source project |
4 | 4 | // |
5 | | -// Copyright (c) 2020 Apple Inc. and the Swift project authors |
| 5 | +// Copyright (c) 2021 Apple Inc. and the Swift project authors |
6 | 6 | // Licensed under Apache License v2.0 with Runtime Library Exception |
7 | 7 | // |
8 | 8 | // See https://swift.org/LICENSE.txt for license information |
@@ -204,3 +204,156 @@ extension Collection { |
204 | 204 | } |
205 | 205 | } |
206 | 206 |
|
| 207 | +//===----------------------------------------------------------------------===// |
| 208 | +// partitioned(_:) |
| 209 | +//===----------------------------------------------------------------------===// |
| 210 | + |
| 211 | +extension Sequence { |
| 212 | + /// Returns two arrays containing, in order, the elements of the sequence that |
| 213 | + /// do and don’t satisfy the given predicate, respectively. |
| 214 | + /// |
| 215 | + /// In this example, `partitioned(_:)` is used to separate the input based on |
| 216 | + /// names that aren’t and are shorter than five characters, respectively: |
| 217 | + /// |
| 218 | + /// let cast = ["Vivien", "Marlon", "Kim", "Karl"] |
| 219 | + /// let (longNames, shortNames) = cast.partitioned({ $0.count < 5 }) |
| 220 | + /// print(longNames) |
| 221 | + /// // Prints "["Vivien", "Marlon"]" |
| 222 | + /// print(shortNames) |
| 223 | + /// // Prints "["Kim", "Karl"]" |
| 224 | + /// |
| 225 | + /// - Parameter belongsInSecondCollection: A closure that takes an element of |
| 226 | + /// the sequence as its argument and returns a Boolean value indicating |
| 227 | + /// whether the element should be included in the second returned array. |
| 228 | + /// Otherwise, the element will appear in the first returned array. |
| 229 | + /// |
| 230 | + /// - Returns: Two arrays with with all of the elements of the receiver. The |
| 231 | + /// first array contains all the elements that `belongsInSecondCollection` |
| 232 | + /// didn’t allow, and the second array contains all the elements that |
| 233 | + /// `belongsInSecondCollection` allowed. |
| 234 | + /// |
| 235 | + /// - Complexity: O(*n*), where *n* is the length of the sequence. |
| 236 | + /// |
| 237 | + /// - Note: This algorithm performs a bit slower than the same algorithm on |
| 238 | + /// `RandomAccessCollection` since the size of the sequence is unknown, unlike |
| 239 | + /// `RandomAccessCollection`. |
| 240 | + @inlinable |
| 241 | + public func partitioned( |
| 242 | + _ belongsInSecondCollection: (Element) throws -> Bool |
| 243 | + ) rethrows -> ([Element], [Element]) { |
| 244 | + var lhs = ContiguousArray<Element>() |
| 245 | + var rhs = ContiguousArray<Element>() |
| 246 | + |
| 247 | + for element in self { |
| 248 | + if try belongsInSecondCollection(element) { |
| 249 | + rhs.append(element) |
| 250 | + } else { |
| 251 | + lhs.append(element) |
| 252 | + } |
| 253 | + } |
| 254 | + |
| 255 | + return _tupleMap((lhs, rhs), { Array($0) }) |
| 256 | + } |
| 257 | +} |
| 258 | + |
| 259 | +extension Collection { |
| 260 | + // This is a specialized version of the same algorithm on `Sequence` that |
| 261 | + // avoids reallocation of arrays since `count` is known ahead of time. |
| 262 | + @inlinable |
| 263 | + public func partitioned( |
| 264 | + _ belongsInSecondCollection: (Element) throws -> Bool |
| 265 | + ) rethrows -> ([Element], [Element]) { |
| 266 | + guard !self.isEmpty else { |
| 267 | + return ([], []) |
| 268 | + } |
| 269 | + |
| 270 | + // Since `RandomAccessCollection`s have known sizes (access to `count` is |
| 271 | + // constant time, O(1)), we can allocate one array of size `self.count`, |
| 272 | + // then insert items at the beginning or end of that contiguous block. This |
| 273 | + // way, we don’t have to do any dynamic array resizing. Since we insert the |
| 274 | + // right elements on the right side in reverse order, we need to reverse |
| 275 | + // them back to the original order at the end. |
| 276 | + |
| 277 | + let count = self.count |
| 278 | + |
| 279 | + // Inside of the `initializer` closure, we set what the actual mid-point is. |
| 280 | + // We will use this to partitioned the single array into two in constant time. |
| 281 | + var midPoint: Int = 0 |
| 282 | + |
| 283 | + let elements = try [Element]( |
| 284 | + unsafeUninitializedCapacity: count, |
| 285 | + initializingWith: { buffer, initializedCount in |
| 286 | + var lhs = buffer.baseAddress! |
| 287 | + var rhs = lhs + buffer.count |
| 288 | + do { |
| 289 | + for element in self { |
| 290 | + if try belongsInSecondCollection(element) { |
| 291 | + rhs -= 1 |
| 292 | + rhs.initialize(to: element) |
| 293 | + } else { |
| 294 | + lhs.initialize(to: element) |
| 295 | + lhs += 1 |
| 296 | + } |
| 297 | + } |
| 298 | + |
| 299 | + let rhsIndex = rhs - buffer.baseAddress! |
| 300 | + buffer[rhsIndex...].reverse() |
| 301 | + initializedCount = buffer.count |
| 302 | + |
| 303 | + midPoint = rhsIndex |
| 304 | + } catch { |
| 305 | + let lhsCount = lhs - buffer.baseAddress! |
| 306 | + let rhsCount = (buffer.baseAddress! + buffer.count) - rhs |
| 307 | + buffer.baseAddress!.deinitialize(count: lhsCount) |
| 308 | + rhs.deinitialize(count: rhsCount) |
| 309 | + throw error |
| 310 | + } |
| 311 | + }) |
| 312 | + |
| 313 | + let collections = elements.partitioned(upTo: midPoint) |
| 314 | + return _tupleMap(collections, { Array($0) }) |
| 315 | + } |
| 316 | +} |
| 317 | + |
| 318 | +//===----------------------------------------------------------------------===// |
| 319 | +// partitioned(upTo:) |
| 320 | +//===----------------------------------------------------------------------===// |
| 321 | + |
| 322 | +extension Collection { |
| 323 | + /// Splits the receiving collection into two at the specified index |
| 324 | + /// - Parameter index: The index within the receiver to split the collection |
| 325 | + /// - Returns: A tuple with the first and second parts of the receiving |
| 326 | + /// collection after splitting it |
| 327 | + /// - Note: The first subsequence in the returned tuple does *not* include |
| 328 | + /// the element at `index`. That element is in the second subsequence. |
| 329 | + /// - Complexity: O(*1*) |
| 330 | + @inlinable |
| 331 | + public func partitioned(upTo index: Index) -> (SubSequence, SubSequence) { |
| 332 | + return ( |
| 333 | + self[self.startIndex..<index], |
| 334 | + self[index..<self.endIndex] |
| 335 | + ) |
| 336 | + } |
| 337 | +} |
| 338 | + |
| 339 | +//===----------------------------------------------------------------------===// |
| 340 | +// _tupleMap(_:_:) |
| 341 | +//===----------------------------------------------------------------------===// |
| 342 | + |
| 343 | +/// Returns a tuple containing the results of mapping the given closure over |
| 344 | +/// each of the tuple’s elements. |
| 345 | +/// - Parameters: |
| 346 | +/// - x: The tuple to transform |
| 347 | +/// - transform: A mapping closure. `transform` accepts an element of this |
| 348 | +/// sequence as its parameter and returns a transformed |
| 349 | +/// - Returns: A tuple containing the transformed elements of this tuple. |
| 350 | +@usableFromInline |
| 351 | +internal func _tupleMap<T, U>( |
| 352 | + _ x: (T, T), |
| 353 | + _ transform: (T) throws -> U |
| 354 | +) rethrows -> (U, U) { |
| 355 | + return ( |
| 356 | + try transform(x.0), |
| 357 | + try transform(x.1) |
| 358 | + ) |
| 359 | +} |
0 commit comments