100日アルゴリズム[18日目・ヒープ]

解いた問題

https://leetcode.com/problems/relative-ranks/description/?envType=problem-list-v2&envId=heap-priority-queue

解法

maxHeap(最大ヒープ)を使ってとく。最大ヒープとは親ノードの値が、子ノードの値より大きいもしくは等しくなっている二分木を指す。逆のものを最小ヒープという。

最小ヒープの実装

配列とそのインデックス番号で考える。

class MinHeap {
    private heap:[] = []

    constructor() {
        this.heap = []
    }

    parentIndex(index:number):number {
        return index % 2;
    }

    leftChildIndex(index:number):number {
        return 2 * index;
    }

    rightChildIndex(index:number):number {
        return (2*index) + 1
    }

    swapElement(index1:number, index2:number) {
        this.heap[index1], this.heap[index2] = this.heap[index2], this.heap[index1]
    }
}

解法

maxHeap(最大ヒープ)を組み合わせて解いたものが以下。

interface ScoreEntry {
    score: number
    idx: number
}

function findRelativeRanks(scores: number[]): string[] {
    const heap = new MaxHeap<ScoreEntry>((a, b) => a.score - b.score)

    for (let i = 0; i < scores.length; i++) {
        heap.insert({ score: scores[i], idx: i })
    }

    const result: string[] = new Array(scores.length)
    const medals = ['Gold Medal', 'Silver Medal', 'Bronze Medal']
    
    for (let i = 0; heap.size() > 0; i++) {
        const { idx } = heap.extractMax()
        result[idx] = medals[i] || String(i + 1)
    }

    return result
};

class MaxHeap<T> {
    heap: T[] = []
    comparator: (a: T, b: T) => number

    constructor(comparator: (a: T, b: T) => number) {
        this.comparator = comparator
    }

    insert(node: T): void {
        let idx = this.heap.push(node) - 1

        while (idx > 0) {
            const parentIdx = this.getParentIdx(idx)
            if (this.comparator(this.heap[idx], this.heap[parentIdx]) > 0) {
                this.swap(idx, parentIdx)
                idx = parentIdx
            } else {
                break
            }
        }
    }

    extractMax(): T | undefined {
        if (this.heap.length === 0) {
            return undefined
        }

        const root = this.heap[0]

        this.heap[0] = this.heap.at(-1)
        this.heap.length -= 1

        let idx = 0
        while (idx < this.heap.length) {
            const leftChildIdx = this.getLeftChildIdx(idx)
            const rightChildIdx = this.getRightChildIdx(idx)

            let largest = idx
            if (
                leftChildIdx < this.heap.length && 
                this.comparator(this.heap[leftChildIdx], this.heap[largest]) > 0
            ) {
                largest = leftChildIdx
            }

            if (
                rightChildIdx < this.heap.length && 
                this.comparator(this.heap[rightChildIdx], this.heap[largest]) > 0
            ) {
                largest = rightChildIdx
            }

            if (largest === idx) break

            this.swap(idx, largest)
            idx = largest
        }

        return root
    }

    size(): number {
        return this.heap.length
    }

    getParentIdx(index: number): number {
        return Math.floor((index - 1) / 2)
    }

    getLeftChildIdx(index: number): number {
        return 2 * index + 1
    }

    getRightChildIdx(index: number): number {
        return 2 * index + 2
    }

    swap(idx1: number, idx2): void {
        [this.heap[idx1], this.heap[idx2]] = [this.heap[idx2], this.heap[idx1]]
    }
}