Find Median from Data Stream

Leetcode

Solution

Approach 1: Using Binary Search Tree

class Node {
    public:
    int val;
    int size;
    Node *left;
    Node *right;
    Node(int v) {
        val = v;
        size = 1;
        left = nullptr;
        right = nullptr;
    }
};

class BinarySearchTree {
    public:
    Node* root = nullptr;
    int size = 0;

    BinarySearchTree() {}

    void addNum(int val) {
        root = addNode(root, val);
        size++;
        // printf("Added: %d\n", root->size);
    }

    int findKthNum(int k) {
        Node* n = findKthNode(root, k);
        if (n != nullptr) {
            return n->val;
        }
        return -1;
    }

    void printTree(Node* root, int indent) {
        if (root == nullptr) {
            return;
        }
        for (int i = 0; i < indent; i++) {
            printf("  ");
        }
        cout << root;
        printf(" %d (Size: %d)\n", root->val, root->size);
        // cout << root << " " << "Size"
        printTree(root->left, indent+1);
        printTree(root->right, indent+1);
    }

    Node* addNode(Node* root, int val) {
        if (root == nullptr) {
            return new Node(val);
        }
        root->size++;
        // printf("Inserting %d\t", val);
        // printf("at node: %d Size: %d\n", root->val, root->size);
        if (root->val >= val) {
            root->left = addNode(root->left, val);
            return root;
        }
        root->right = addNode(root->right, val);
        return root;
    }

    Node* findKthNode(Node* root, int k) {
        if (root == nullptr) { // assume that we never reach here
            return nullptr;
        }

        // printf("Finding %d at ", k);
        // cout << root << "\n";
        int leftSize = 0;
        if (root->left != nullptr) {
            leftSize = root->left->size;
        }
        if (k <= leftSize) {
            return findKthNode(root->left, k);
        } else if (k == leftSize + 1) {
            return root;
        } else {
            return findKthNode(root->right, k - 1 - leftSize);
        }
    }
};

class MedianFinder {
public:
    BinarySearchTree* bt = new BinarySearchTree();
    MedianFinder() {}

    void addNum(int num) {
        bt->addNum(num);

        // printf("Added: %d\n", num);
    }
    
    double findMedian() {
        int size = bt->size;
        // if (size > 3) {
        //     bt->printTree(bt->root, 0);
        // }
        if (size % 2) {
            int k = (size / 2) + 1;
            return bt->findKthNum(k);
        } else {
            int k1 = (size / 2);
            int k2 = (size / 2) + 1;
            int d1 = bt->findKthNum(k1);
            int d2 = bt->findKthNum(k2);
            // printf("Find: %d, %d\n", k1, d1);
            // printf("Find: %d, %d\n", k2, d2);

            return (double(d1)+double(d2)) / double(2);
        }
    }
};

/**
 * Your MedianFinder object will be instantiated and called as such:
 * MedianFinder* obj = new MedianFinder();
 * obj->addNum(num);
 * double param_2 = obj->findMedian();
*/

The above solution will time out since the binary search tree is not balanced.

Approach 2: Using Two Heaps

class MedianFinder {
public:
    priority_queue<int> maxHeap;
    priority_queue<int, vector<int>, greater<int>> minHeap;
    
    MedianFinder() {
    }
    void addNum(int num) {
        maxHeap.push(num);
        minHeap.push(maxHeap.top());
        maxHeap.pop();
        if (minHeap.size() > maxHeap.size()) {
            maxHeap.push(minHeap.top());
            minHeap.pop();
        }
    }
    double findMedian() {
        if (maxHeap.size() > minHeap.size()) return maxHeap.top();
        return (maxHeap.top() + minHeap.top()) / 2.0;
    }
};