Find Median from Data Stream

The median is the middle value in an ordered integer list. If the size of the list is even, there is no middle value and the median is the mean of the two middle values.

  • For example, for arr = [2,3,4], the median is 3.
  • For example, for arr = [2,3], the median is (2 + 3) / 2 = 2.5.

Implement the MedianFinder class:

  • MedianFinder() initializes the MedianFinder object.
  • void addNum(int num) adds the integer num from the data stream to the data structure.
  • double findMedian() returns the median of all elements so far. Answers within 10-5 of the actual answer will be accepted.

Example 1:

Input
["MedianFinder", "addNum", "addNum", "findMedian", "addNum", "findMedian"]
[[], [1], [2], [], [3], []]
Output

[null, null, null, 1.5, null, 2.0]

Explanation MedianFinder medianFinder = new MedianFinder(); medianFinder.addNum(1); // arr = [1] medianFinder.addNum(2); // arr = [1, 2] medianFinder.findMedian(); // return 1.5 (i.e., (1 + 2) / 2) medianFinder.addNum(3); // arr[1, 2, 3] medianFinder.findMedian(); // return 2.0

Constraints:

  • -105 <= num <= 105
  • There will be at least one element in the data structure before calling findMedian.
  • At most 5 * 104 calls will be made to addNum and findMedian.

Follow up:

  • If all integer numbers from the stream are in the range [0, 100], how would you optimize your solution?
  • If 99% of all integer numbers from the stream are in the range [0, 100], how would you optimize your solution?
struct Heap<Element>
{
	var elements : [Element]
	let priorityFunction : (Element, Element) -> Bool

	init(elements: [Element] = [], priorityFunction: @escaping (Element, Element) -> Bool) {
		self.elements = elements
		self.priorityFunction = priorityFunction
		buildHeap()
	}

	mutating func buildHeap() {
		for index in (0 ..< count / 2).reversed() {
			siftDown(elementAtIndex: index)
		}
	}
	
	var isEmpty : Bool { return elements.isEmpty }
	var count : Int { return elements.count }

	func peek() -> Element? {
		return elements.first
	}

	mutating func enqueue(_ element: Element) {
		elements.append(element)
		siftUp(elementAtIndex: count - 1)
	}

	mutating func siftUp(elementAtIndex index: Int) {
		let parent = parentIndex(of: index)
		guard !isRoot(index),
			isHigherPriority(at: index, than: parent)
			else { return }
		swapElement(at: index, with: parent)
		siftUp(elementAtIndex: parent)
	}

	mutating func dequeue() -> Element? {
		guard !isEmpty
			else { return nil }
		swapElement(at: 0, with: count - 1)
		let element = elements.removeLast()
		if !isEmpty {
			siftDown(elementAtIndex: 0)
		}
		return element
	}
	
	mutating func siftDown(elementAtIndex index: Int) {
		let childIndex = highestPriorityIndex(for: index)
		if index == childIndex {
			return
		}
		swapElement(at: index, with: childIndex)
		siftDown(elementAtIndex: childIndex)
	}

	// Helper functions
	
	func isRoot(_ index: Int) -> Bool {
		return index == 0
	}
	
	func leftChildIndex(of index: Int) -> Int {
		return (2 * index) + 1
	}
	
	func rightChildIndex(of index: Int) -> Int {
		return (2 * index) + 2
	}
	
	func parentIndex(of index: Int) -> Int {
		return (index - 1) / 2
	}
	
	func isHigherPriority(at firstIndex: Int, than secondIndex: Int) -> Bool {
		return priorityFunction(elements[firstIndex], elements[secondIndex])
	}
	
	func highestPriorityIndex(of parentIndex: Int, and childIndex: Int) -> Int {
		guard childIndex < count && isHigherPriority(at: childIndex, than: parentIndex)
			else { return parentIndex }
		return childIndex
	}
	
	func highestPriorityIndex(for parent: Int) -> Int {
		return highestPriorityIndex(of: highestPriorityIndex(of: parent, and: leftChildIndex(of: parent)), and: rightChildIndex(of: parent))
	}
	
	mutating func swapElement(at firstIndex: Int, with secondIndex: Int) {
		guard firstIndex != secondIndex
			else { return }
		elements.swapAt(firstIndex, secondIndex)
	}
}

// A bonus extra for you: two extra functions, if the Element type is Equatable
extension Heap where Element : Equatable {
	
	// This function allows you to remove an element from the heap, in a similar way to how you would dequeue the root element.
	mutating func remove(_ element: Element) {
		guard let index = elements.index(of: element)
			else { return }
		swapElement(at: index, with: count - 1)
		elements.remove(at: count - 1)
		siftDown(elementAtIndex: index)
	}

	// This function allows you to 'boost' an element, by sifting the element up the heap. You might do this if the element is already in the heap, but its priority has increased since it was enqueued.
	mutating func boost(_ element: Element) {
		guard let index = elements.index(of: element)
			else { return }
		siftUp(elementAtIndex: index)
	}
}

var heap = Heap<Int>(elements: [3, 2, 8, 5, 0], priorityFunction: >)
heap.enqueue(6)
heap.enqueue(1)
heap.enqueue(4)
heap.enqueue(7)
//heap.dequeue()
//print(heap.elements)


class MedianFinder {
	var modifiedArr = [Double]()
	var minHeap = Heap<Int>(elements: [Int](), priorityFunction: <)
	var maxHeap = Heap<Int>(elements: [Int](), priorityFunction: >)
	init() {
		
	}
	
	func addNum(_ num: Int) {
		modifiedArr.append(Double(num))
	}
	
	func findMedian() -> Double {
		if modifiedArr.count == 1 {
			return modifiedArr[0]
		} else if modifiedArr.count == 2 {
			return Double(modifiedArr[0]+modifiedArr[1])/2.0
		} else {
			modifiedArr.sort()
			if modifiedArr.count % 2 == 0 {
				// even
				let middle = modifiedArr.count / 2
				return (modifiedArr[middle-1]+modifiedArr[middle])/2
			} else {
				// odd
				let middle = modifiedArr.count / 2
				return modifiedArr[middle]
			}
		}
	}
	
	// Using minHeap and maxHeap
	// https://www.youtube.com/watch?v=1LkOrc-Le-Y
	func addNum2(_ num: Int) {
		// maxHeap is empty
		if maxHeap.isEmpty {
			maxHeap.enqueue(num)
		}
		// heap size is same
		else if minHeap.count == maxHeap.count {
			if num < minHeap.peek()! {
				maxHeap.enqueue(num)
			} else {
				let lowest = minHeap.dequeue()
				minHeap.enqueue(num)
				maxHeap.enqueue(lowest!)
			}
		}
		// unequal heap sizes
		else {
			if minHeap.count == 0 {
				if num > maxHeap.peek()! {
					minHeap.enqueue(num)
				} else {
					let highest = maxHeap.dequeue()!
					maxHeap.enqueue(num)
					minHeap.enqueue(highest)
				}
			} else if num >= minHeap.peek()! {
				minHeap.enqueue(num)
			} else {
				if num < maxHeap.peek()! {
					let highest = maxHeap.peek()!
					maxHeap.dequeue()
					maxHeap.enqueue(num)
					minHeap.enqueue(highest)
				} else {
					minHeap.enqueue(num)
				}
			}
		}
	}
	
	func findMedian2() -> Double {
		if maxHeap.count > minHeap.count {
			return Double(maxHeap.peek()!)
		} else {
			print(minHeap.peek()!, maxHeap.peek()!)
			return Double(((Double(minHeap.peek()!) + Double(maxHeap.peek()!)) / 2.0))
		}
	}
}

var medianFinder = MedianFinder()
// medianFinder.addNum(5);    // arr = [1]
// medianFinder.addNum(2);    // arr = [1, 2]
// print(medianFinder.findMedian()) // return 1.5 (i.e., (1 + 2) / 2)
// medianFinder.addNum(3);    // arr[1, 2, 3]
// medianFinder.findMedian(); // return 2.0
medianFinder.addNum2(5)
print(medianFinder.findMedian2())
medianFinder.addNum2(3)
print(medianFinder.findMedian2())
medianFinder.addNum2(4)
print(medianFinder.findMedian2())
medianFinder.addNum2(2)
print(medianFinder.findMedian2())
medianFinder.addNum2(6)
print(medianFinder.findMedian2())

Leave a Reply

Your email address will not be published.