1 class Solution {
  2     func mincostToHireWorkers(_ quality: [Int], _ wage: [Int], _ K: Int) -> Double {
  3   var workers: [(wage: Int, quality: Int)] = []
  4         
  5         for i in 0..<quality.count {
  6             workers.append((wage[i], quality[i]))
  7         }
  8         // Sort all workers with their expected ratio
  9         workers.sort { Double($0.wage) / Double($0.quality) < Double($1.wage) / Double($1.quality) } // O(n*logn)
 10         
 11         var heap = Heap(sort: >)
 12         
 13         var totalWage = Double(Int.max)
 14         var totalQuality = 0
 15         
 16         for worker in workers {
 17             let currRatio = Double(worker.wage)/Double(worker.quality)
 18             //print(currRatio)
 19             heap.push(worker.quality)
 20             totalQuality += worker.quality
 21             //print("total quality:", totalQuality)
 22             if heap.count > K {
 23                 totalQuality -= heap.remove()!
 24             }
 25             if heap.count == K {
 26                 totalWage = min(totalWage, Double(totalQuality) * currRatio)
 27                 //print("current total wage:", totalWage)
 28             }
 29         }
 30         
 31         return totalWage
 32     }
 33 }
 34 
 35 // Implementation of heap
 36 public struct Heap {
 37     
 38     var elements: [Int] = []
 39     let sort: (Int, Int) -> Bool
 40     var isEmpty: Bool {
 41         return self.elements.isEmpty
 42     }
 43     
 44     var count: Int {
 45         return self.elements.count
 46     }
 47     
 48     func peek() -> Int? {
 49         return elements.first
 50     }
 51     
 52     init(sort: @escaping (Int, Int) -> Bool, elements: [Int] = []) {
 53         self.sort = sort
 54         self.elements = elements
 55         
 56         if !elements.isEmpty {
 57             for i in stride(from: elements.count/2 - 1, through: 0, by: -1) {
 58                 siftDown(from: i)
 59             }
 60         }
 61     }
 62     
 63     mutating func siftDown(from index: Int) {
 64         var parent = index
 65         while true {
 66             let left = leftIndex(of: parent)
 67             let right = rightIndex(of: parent)
 68             var candidate = parent
 69             if left < count && sort(elements[left], elements[candidate]) {
 70                 candidate = left
 71             }
 72             if right < count && sort(elements[right], elements[candidate]) {
 73                 candidate = right
 74             }
 75             if candidate == parent {
 76                 return
 77             }
 78             elements.swapAt(parent, candidate)
 79             parent = candidate
 80         }
 81     }
 82     
 83     mutating func siftUp(from index: Int) {
 84         var child = index
 85         var parent = parentIndex(of: child)
 86         while child > 0 && sort(elements[child], elements[parent]) {
 87             elements.swapAt(child, parent)
 88             child = parent
 89             parent = parentIndex(of: child)
 90         }
 91     }
 92     
 93     mutating func push(_ element: Int) {
 94         elements.append(element)
 95         siftUp(from: count-1)
 96     }
 97     
 98     mutating func remove() -> Int? {
 99         guard !isEmpty else { return nil }
100         elements.swapAt(0, count-1)
101         defer {
102             siftDown(from: 0)
103         }
104         return elements.popLast()
105     }
106     
107     func leftIndex(of index: Int) -> Int {
108         return (2 * index) + 1
109     }
110     
111     func rightIndex(of index: Int) -> Int {
112         return (2 * index) + 2
113     }
114     
115     func parentIndex(of index: Int) -> Int {
116         return (index - 1) / 2
117     }
118 }

  1 class Solution {
  2     func mincostToHireWorkers(_ quality: [Int], _ wage: [Int], _ K: Int) -> Double {
  3         var quaRatios = [(Int, Double)]()
  4         for i in 0..<quality.count {
  5             quaRatios.append((quality[i], Double(wage[i]) / Double(quality[i])))
  6         }
  7         quaRatios.sort { $0.1 < $1.1 }
  8         var pq = PriorityQueue<Int> { $0 > $1 }
  9         var sum = 0
 10         var result = Double(Int.max)
 11         for (qua, ratio) in quaRatios {
 12             sum += qua
 13             pq.enqueue(qua)
 14             if pq.count > K {
 15                 sum -= pq.dequeue()!
 16             }
 17             if pq.count == K {
 18                 result = min(result, ratio * Double(sum))
 19             }
 20         }
 21         return result
 22     }
 23 }
 24 
 25 public struct PriorityQueue<T> {
 26   fileprivate var heap: Heap<T>
 27 
 28   /*
 29     To create a max-priority queue, supply a > sort function. For a min-priority
 30     queue, use <.
 31   */
 32   public init(sort: @escaping (T, T) -> Bool) {
 33     heap = Heap(sort: sort)
 34   }
 35 
 36   public var isEmpty: Bool {
 37     return heap.isEmpty
 38   }
 39 
 40   public var count: Int {
 41     return heap.count
 42   }
 43 
 44   public func peek() -> T? {
 45     return heap.peek()
 46   }
 47 
 48   public mutating func enqueue(_ element: T) {
 49     heap.insert(element)
 50   }
 51 
 52   public mutating func dequeue() -> T? {
 53     return heap.remove()
 54   }
 55 
 56   /*
 57     Allows you to change the priority of an element. In a max-priority queue,
 58     the new priority should be larger than the old one; in a min-priority queue
 59     it should be smaller.
 60   */
 61   public mutating func changePriority(index i: Int, value: T) {
 62     return heap.replace(index: i, value: value)
 63   }
 64 }
 65 
 66 extension PriorityQueue where T: Equatable {
 67   public func index(of element: T) -> Int? {
 68     return heap.index(of: element)
 69   }
 70 }
 71 
 72 //
 73 //  Heap.swift
 74 //  Written for the Swift Algorithm Club by Kevin Randrup and Matthijs Hollemans
 75 //
 76 public struct Heap<T> {
 77   
 78   /** The array that stores the heap's nodes. */
 79   var nodes = [T]()
 80   
 81   /**
 82    * Determines how to compare two nodes in the heap.
 83    * Use '>' for a max-heap or '<' for a min-heap,
 84    * or provide a comparing method if the heap is made
 85    * of custom elements, for example tuples.
 86    */
 87   private var orderCriteria: (T, T) -> Bool
 88   
 89   /**
 90    * Creates an empty heap.
 91    * The sort function determines whether this is a min-heap or max-heap.
 92    * For comparable data types, > makes a max-heap, < makes a min-heap.
 93    */
 94   public init(sort: @escaping (T, T) -> Bool) {
 95     self.orderCriteria = sort
 96   }
 97   
 98   /**
 99    * Creates a heap from an array. The order of the array does not matter;
100    * the elements are inserted into the heap in the order determined by the
101    * sort function. For comparable data types, '>' makes a max-heap,
102    * '<' makes a min-heap.
103    */
104   public init(array: [T], sort: @escaping (T, T) -> Bool) {
105     self.orderCriteria = sort
106     configureHeap(from: array)
107   }
108   
109   /**
110    * Configures the max-heap or min-heap from an array, in a bottom-up manner.
111    * Performance: This runs pretty much in O(n).
112    */
113   private mutating func configureHeap(from array: [T]) {
114     nodes = array
115     for i in stride(from: (nodes.count/2-1), through: 0, by: -1) {
116       shiftDown(i)
117     }
118   }
119   
120   public var isEmpty: Bool {
121     return nodes.isEmpty
122   }
123   
124   public var count: Int {
125     return nodes.count
126   }
127   
128   /**
129    * Returns the index of the parent of the element at index i.
130    * The element at index 0 is the root of the tree and has no parent.
131    */
132   @inline(__always) internal func parentIndex(ofIndex i: Int) -> Int {
133     return (i - 1) / 2
134   }
135   
136   /**
137    * Returns the index of the left child of the element at index i.
138    * Note that this index can be greater than the heap size, in which case
139    * there is no left child.
140    */
141   @inline(__always) internal func leftChildIndex(ofIndex i: Int) -> Int {
142     return 2*i + 1
143   }
144   
145   /**
146    * Returns the index of the right child of the element at index i.
147    * Note that this index can be greater than the heap size, in which case
148    * there is no right child.
149    */
150   @inline(__always) internal func rightChildIndex(ofIndex i: Int) -> Int {
151     return 2*i + 2
152   }
153   
154   /**
155    * Returns the maximum value in the heap (for a max-heap) or the minimum
156    * value (for a min-heap).
157    */
158   public func peek() -> T? {
159     return nodes.first
160   }
161   
162   /**
163    * Adds a new value to the heap. This reorders the heap so that the max-heap
164    * or min-heap property still holds. Performance: O(log n).
165    */
166   public mutating func insert(_ value: T) {
167     nodes.append(value)
168     shiftUp(nodes.count - 1)
169   }
170   
171   /**
172    * Adds a sequence of values to the heap. This reorders the heap so that
173    * the max-heap or min-heap property still holds. Performance: O(log n).
174    */
175   public mutating func insert<S: Sequence>(_ sequence: S) where S.Iterator.Element == T {
176     for value in sequence {
177       insert(value)
178     }
179   }
180   
181   /**
182    * Allows you to change an element. This reorders the heap so that
183    * the max-heap or min-heap property still holds.
184    */
185   public mutating func replace(index i: Int, value: T) {
186     guard i < nodes.count else { return }
187     
188     remove(at: i)
189     insert(value)
190   }
191   
192   /**
193    * Removes the root node from the heap. For a max-heap, this is the maximum
194    * value; for a min-heap it is the minimum value. Performance: O(log n).
195    */
196   @discardableResult public mutating func remove() -> T? {
197     guard !nodes.isEmpty else { return nil }
198     
199     if nodes.count == 1 {
200       return nodes.removeLast()
201     } else {
202       // Use the last node to replace the first one, then fix the heap by
203       // shifting this new first node into its proper position.
204       let value = nodes[0]
205       nodes[0] = nodes.removeLast()
206       shiftDown(0)
207       return value
208     }
209   }
210   
211   /**
212    * Removes an arbitrary node from the heap. Performance: O(log n).
213    * Note that you need to know the node's index.
214    */
215   @discardableResult public mutating func remove(at index: Int) -> T? {
216     guard index < nodes.count else { return nil }
217     
218     let size = nodes.count - 1
219     if index != size {
220       nodes.swapAt(index, size)
221       shiftDown(from: index, until: size)
222       shiftUp(index)
223     }
224     return nodes.removeLast()
225   }
226   
227   /**
228    * Takes a child node and looks at its parents; if a parent is not larger
229    * (max-heap) or not smaller (min-heap) than the child, we exchange them.
230    */
231   internal mutating func shiftUp(_ index: Int) {
232     var childIndex = index
233     let child = nodes[childIndex]
234     var parentIndex = self.parentIndex(ofIndex: childIndex)
235     
236     while childIndex > 0 && orderCriteria(child, nodes[parentIndex]) {
237       nodes[childIndex] = nodes[parentIndex]
238       childIndex = parentIndex
239       parentIndex = self.parentIndex(ofIndex: childIndex)
240     }
241     
242     nodes[childIndex] = child
243   }
244   
245   /**
246    * Looks at a parent node and makes sure it is still larger (max-heap) or
247    * smaller (min-heap) than its childeren.
248    */
249   internal mutating func shiftDown(from index: Int, until endIndex: Int) {
250     let leftChildIndex = self.leftChildIndex(ofIndex: index)
251     let rightChildIndex = leftChildIndex + 1
252     
253     // Figure out which comes first if we order them by the sort function:
254     // the parent, the left child, or the right child. If the parent comes
255     // first, we're done. If not, that element is out-of-place and we make
256     // it "float down" the tree until the heap property is restored.
257     var first = index
258     if leftChildIndex < endIndex && orderCriteria(nodes[leftChildIndex], nodes[first]) {
259       first = leftChildIndex
260     }
261     if rightChildIndex < endIndex && orderCriteria(nodes[rightChildIndex], nodes[first]) {
262       first = rightChildIndex
263     }
264     if first == index { return }
265     
266     nodes.swapAt(index, first)
267     shiftDown(from: first, until: endIndex)
268   }
269   
270   internal mutating func shiftDown(_ index: Int) {
271     shiftDown(from: index, until: nodes.count)
272   }
273   
274 }
275 
276 // MARK: - Searching
277 extension Heap where T: Equatable {
278   
279   /** Get the index of a node in the heap. Performance: O(n). */
280   public func index(of node: T) -> Int? {
281     return nodes.index(where: { $0 == node })
282   }
283   
284   /** Removes the first occurrence of a node from the heap. Performance: O(n log n). */
285   @discardableResult public mutating func remove(node: T) -> T? {
286     if let index = index(of: node) {
287       return remove(at: index)
288     }
289     return nil
290   }  
291 }