算法笔记38:合并K个升序链表

83 阅读1分钟

23. 合并K个升序链表

这道题目除了可以用二分法降低复杂度之外,还可以使用优先队列。主要思想就是把每个链表的头先都放到优先队列里,这样往外取的时候就可以按升序取了。

每取出来一个,就从取出来的这个往后去找,如果还有节点,就再加到队列里,直到队列为空。

class PriorityQueue {
    queue = [];
    
    constructor(comparator) {
        // allow user to pass in the custom comparator
        this.comparator = comparator || function(a, b) { return a - b; };
    }

    parentIndex(i) {
        return Math.floor((i - 1) / 2)
    }
    
    leftChildIndex(i) {
        return i * 2 + 1;
    }
    
    rightChildIndex(i) {
        return i * 2 + 2;
    }

    peek() {
        return this.queue[0];
    }

    size() {
        return this.queue.length;
    }
  
    offer(item) {
        this.queue.push(item);
        // when add new item in, shift the last one up to correct position
        this.shiftUp(this.queue.length - 1);
    }
    
    poll() {
        const item = this.queue.shift();

        if (this.queue.length !== 0) {
            const lastLeaf = this.queue.pop();
            this.queue.unshift(lastLeaf);
            // shift the last one down
            this.shiftDown(0);
        }    
        return item;
    }
    
    shiftUp(i) {
        const parentIndex = this.parentIndex(i);
        // compare with its parent
        if (i > 0 && this.comparator(this.queue[parentIndex], this.queue[i]) < 0) {
            // if parent is smaller, swap them
            this.swap(i, parentIndex);
            // continue to shift up
            this.shiftUp(parentIndex);
        }
    }
    
    shiftDown(i) {
        let largest = i;
        const l = this.leftChildIndex(i);
        const r = this.rightChildIndex(i);
        
        // figure out which one is the largest among parent and two children
        if (l < this.queue.length &&
            this.comparator(this.queue[largest], this.queue[l]) < 0) {
            largest = l;
        }
        if (r < this.queue.length &&
            this.comparator(this.queue[largest], this.queue[r]) < 0) {
            largest = r;
        }

        if (largest != i) {
            // if parent is not the largest, shift down
            this.swap(i, largest);
            // continue to shift down
            this.shiftDown(largest);
        }
    }
    
    swap(l, r) {
        const t = this.queue[l];
        this.queue[l] = this.queue[r];
        this.queue[r] = t;
    }
}

const mergeKLists = function(lists) {
    const pQueue = new PriorityQueue((a, b) => b.val - a.val);

    const tmp = new ListNode(null);
    for (const node of lists) {
        node && pQueue.offer(node);
    }

    let p = tmp;
    while (pQueue.size() > 0) {
        const nextNode = pQueue.poll();
        if (nextNode.next) {
            pQueue.offer(nextNode.next);
        }

        p.next = nextNode;
        p = p.next;
    }

    return tmp.next;
};