Improve runtime of spell function with dictionaries (#126)

dn-m · Oct 21, 2018 · 76047b6 · 76047b6
1 parent f86047b
commit 76047b6
Show file tree

Hide file tree

Showing 3 changed files with 20 additions and 14 deletions.
diff --git a/Sources/SpelledPitch/PitchSpeller/Wetherfield/Wetherfield.swift b/Sources/SpelledPitch/PitchSpeller/Wetherfield/Wetherfield.swift
@@ -48,8 +48,8 @@ struct PitchSpeller {
 
     /// - Returns: An array of nodes, each representing the index of the unassigned node in
     /// `pitchNodes`.
-    private static func internalNodes(pitches: [Pitch]) -> [Cross<Int, Tendency>] {
-        return pitches.indices.flatMap { offset in [.down,.up].map { index in node(offset, index) } }
+    private static func internalNodes(pitches: [Int: Pitch]) -> [Cross<Int, Tendency>] {
+        return pitches.keys.flatMap { offset in [.down,.up].map { index in node(offset, index) } }
     }
 
     // MARK: - Instance Properties
@@ -58,7 +58,7 @@ struct PitchSpeller {
     let parsimonyPivot: Pitch.Spelling
 
     /// The unspelled `Pitch` values to be spelled.
-    let pitches: [Pitch]
+    let pitches: [Int: Pitch]
 
     /// The nodes within the `FlowNetwork`. The values are the encodings of the indices of `Pitch`
     /// values in `pitches.
@@ -70,7 +70,7 @@ struct PitchSpeller {
     // MARK: - Initializers
 
     /// Create a `PitchSpeller` to spell the given `pitches`, with the given `parsimonyPivot`.
-    init(pitches: [Pitch], parsimonyPivot: Pitch.Spelling = .init(.d)) {
+    init(pitches: [Int: Pitch], parsimonyPivot: Pitch.Spelling = .init(.d)) {
         self.pitches = pitches
         self.parsimonyPivot = parsimonyPivot
         self.pitchNodes = PitchSpeller.internalNodes(pitches: pitches)
@@ -81,9 +81,9 @@ struct PitchSpeller {
         )
     }
 
-    /// - Returns: An array of `SpelledPitch` values in the order in which the original
-    /// unspelled `Pitch` values are given.
-    func spell() -> [SpelledPitch] {
+    /// - Returns: An array of `SpelledPitch` values with the same indices as the original
+    /// unspelled `Pitch` values.
+    func spell() -> [Int: SpelledPitch] {
 
         var assignedNodes: [AssignedNode] {
             var (sourceSide, sinkSide) = flowNetwork.minimumCut
@@ -95,9 +95,15 @@ struct PitchSpeller {
         }
 
         return assignedNodes
-            .sorted()
-            .pairs
-            .map(spellPitch)
+            .reduce(into: [Int: (AssignedNode, AssignedNode)]()) { pairs, node in
+                if !pairs.keys.contains(node.index.a) {
+                    pairs[node.index.a] = (node, node)
+                }
+                switch node.index.b {
+                case .up: pairs[node.index.a]!.0 = node
+                case .down: pairs[node.index.a]!.1 = node
+                }
+            }.mapValues(spellPitch)
     }
 
     private func spellPitch(_ up: AssignedNode, _ down: AssignedNode) -> SpelledPitch {
@@ -109,7 +115,7 @@ struct PitchSpeller {
 
     /// - Returns: The `Pitch` value for the given `node` value.
     private func pitch(node: Cross<Int, Tendency>) -> Pitch {
-        return pitches[node.a]
+        return pitches[node.a]!
     }
 }
 

diff --git a/Tests/SpelledPitchTests/PitchSpellerTests.swift b/Tests/SpelledPitchTests/PitchSpellerTests.swift
@@ -20,7 +20,7 @@ class PitchSpellerTests: XCTestCase {
     }
 
     func testSpellCF() {
-        let pitchSpeller = PitchSpeller(pitches: [60,65])
+        let pitchSpeller = PitchSpeller(pitches: [0:60,1:65])
         let _ = pitchSpeller.spell()
         // FIXME: Add assertion
     }

diff --git a/Tests/SpelledPitchTests/WetherfieldTests.swift b/Tests/SpelledPitchTests/WetherfieldTests.swift
@@ -13,12 +13,12 @@ import Pitch
 class WetherfieldTests: XCTestCase {
 
     func testInitMonadNodeCount() {
-        let speller = PitchSpeller(pitches: [60], parsimonyPivot: Pitch.Spelling(.d))
+        let speller = PitchSpeller(pitches: [0:60], parsimonyPivot: Pitch.Spelling(.d))
         XCTAssertEqual(speller.flowNetwork.internalNodes.count, 2)
     }
 
     func testInitDyadNodeCount() {
-        let speller = PitchSpeller(pitches: [60,61], parsimonyPivot: Pitch.Spelling(.d))
+        let speller = PitchSpeller(pitches: [0:60,1:61], parsimonyPivot: Pitch.Spelling(.d))
         XCTAssertEqual(speller.flowNetwork.internalNodes.count, 4)
     }
 }