TLE with hand-written heap, can anyone help me identify the problem?


  • 0
    V

    I use my own version of heap just to practice. nRemain is the number of non-empty lists. It is time-limit-exceeded at the input consisting of all singleton lists.

    class Solution {
    public:
        ListNode* mergeKLists(vector<ListNode*>& lists) {
    		int nRemain = lists.size(), i = 0;
    		ListNode head(0), *cptr = &head;
    
    		while (i < nRemain) {
    			if (!lists[i]) { swap(lists[i], lists[nRemain-1]); --nRemain; }
    			else ++i;
    		}
    		for (i = (nRemain-1)/2; i >= 0; --i) heapify(lists, i, nRemain);
    		while (nRemain > 1) {
    			cptr->next = lists[0];
    			cptr = cptr->next;
    			lists[0] = lists[0]->next;
    			if (!lists[0]) { lists[0] = lists[nRemain-1]; --nRemain; }
    			if (nRemain > 1) heapify(lists, 0, nRemain);
    		}		
    		
    		return head.next;
        }
    
    private:
    	void heapify(vector<ListNode*> lists, int i, int& size) {
    		int l = i*2 + 1, r = i*2 + 2, smallest = i;
    		if (l < size && lists[l]->val<lists[i]->val) smallest = l;
    		if (r < size && lists[r]->val<lists[smallest]->val) smallest = r;
    		if (smallest != i) {
    			swap(lists[i], lists[smallest]);
    			heapify(lists, smallest, size);
    		}
    	}
    };
    

    Thanks!


Log in to reply
 

Looks like your connection to LeetCode Discuss was lost, please wait while we try to reconnect.