ヒープと優先度付きキュー

競技プログラミングで度々必要になるヒープをまとめてみる。

ヒープとは

ヒープは木構造の1つで、二分木として表現される。

各ノードがその子ノードより小さい（or 大きい）か等しくなるように配置される。

そのため根ノードが最小（or 最大）になる。

このノードの関係によって、最小ヒープと最大ヒープの2種類がある。

最小ヒープの例

index=1のノードnode(i)と各ノードの関係は以下のようになる。

• 根ノード｜i = 1（i = 0 の実装もあるのでその場合は j = i - 1 として ノードj について考える）

• 親ノード｜parent(i) = i / 2 （小数は切り捨て）

• 左の子ノード｜left(i) = 2i

• 右の子ノード｜right(i) = 2i + 1

例えば0008は配列のindex=5に位置（node(5)）している。

この0008のノードの親ノードは0003のノードでparent(5) = 5 /2 = 2である。

左側の子ノードは0017のノードでleft(5) = 2*5 = 10である。

右側の子ノードは0019のノードでright(5) = 2*5 + 1 = 11である。

競技プログラミングでのヒープ

競技プログラミングでもヒープは度々必要になる。

特に最大値や最小値を取り出して都度何らかの操作をするパターンの問題ではヒープを便利に使うことができる。

それ以外にも、単純にスライスのソートをsortパッケージのsort.Ints()でやっていてTLEしてしまい、初めからスライスに格納するのではなくヒープに格納して最大値や最小値を求める方が早い場合がある。

数字以外にも属性が存在する場合は優先度付きキューを使う。

Goではcontainer/heapパッケージでヒープや優先度付きキューが提供されている。

pkg.go.dev

ヒープで実装する問題

以下の問題を考える。

atcoder.jp

この問題では最大値を求めて取り出し、割引をするという操作をクーポンの数だけ繰り返す必要がある。

この問題を単純に考えて、配列をfor文で回して最大値を求めるというやり方でやると、計算量はO(NM)かかる（1≤N,M≤10⁵）のでTLEしてしまう。

そこでヒープを使って以下のように実装する。

package main

import (
    "bufio"
    "container/heap"
    "fmt"
    "os"
    "strconv"
)

// An IntHeap is a min-heap of ints.
type IntHeap []int

func (h IntHeap) Len() int           { return len(h) }
func (h IntHeap) Less(i, j int) bool { return h[i] > h[j] }
func (h IntHeap) Swap(i, j int)      { h[i], h[j] = h[j], h[i] }

func (h *IntHeap) Push(x interface{}) {
    // Push and Pop use pointer receivers because they modify the slice's length,
    // not just its contents.
    *h = append(*h, x.(int))
}

func (h *IntHeap) Pop() interface{} {
    old := *h
    n := len(old)
    x := old[n-1]
    *h = old[0 : n-1]
    return x
}

var scanner = bufio.NewScanner(os.Stdin)

func nextInt() int {
    scanner.Scan()                       // トークンの読み込み
    i, e := strconv.Atoi(scanner.Text()) // 読み込んだトークンをTextメソッドで取り出し
    if e != nil {
        panic(e)
    }
    return i
}

func main() {
    scanner.Split(bufio.ScanWords) // 空白区切りで受け取るので引数にbufio.ScanwWordsを渡す
    N := nextInt()
    M := nextInt()
    h := &IntHeap{}
    heap.Init(h)
    for i := 0; i < N; i++ {
        heap.Push(h, nextInt())
    }

    for i := 0; i < M; i++ {
        maxValue := heap.Pop(h).(int)
        discountedValue := maxValue / 2
        heap.Push(h, discountedValue)
    }

    var sum int
    for h.Len() > 0 {
        sum += heap.Pop(h).(int)
    }
    fmt.Println(sum)
}