ヒープの強化【競技プログラミング】

はじめに

• 過去、swiftに存在しない構造体であるヒープを自作したことがあった。
• このヒープにremove関数を付けてなかったので、追加することとする。
• また、chgNodeっていう、1カ所のノードの値を更新した場合の関数をおまけで付けたけど、これって、下方向のみ対応していたから、上方向にも対応させたよ。

ヒープへのremove関数追加のロジック

• この時のヒープは、remove関数を装備してなかったけど、今回追加するロジックは以下のとおり。
  1. 取り除きたい値の格納されている配列インデックスi(ノード番号を一般的な1スタートとすれば、ノード番号=インデックス+1)を得る。
  2. 当該ノードの値を、配列の最後尾の値vで更新する。⇒chgNode(v,i)：計算量(O(log N))
  3. 配列の最後尾を削除する。⇒removeLast()!：計算量(O(1))
• 配列を使って実装したため、もし、値指定でremoveを実装すれば、上記1.では、firstIndexでインデックスを得ることになるけど、下記の通り低速。
  • firstIndex(of:)メソッド:O(N)
    • 理由: firstIndex(of:)メソッドは配列の先頭から順に要素を比較していくため、最悪の場合、配列全体を1回走査する必要があります。このため、計算量はO(N)となります。

ヒープのremove高速化のための手当

• 配列の検索関数であるfirstIndexは低速(O(N))なので、コレを高速化したい。
• どうすれば良いか？最近は面倒くさがりなので、すぐに生成AIに聞いちゃう。
  • 自分で調べる場合でもgoogle先生に聞くだけなので、AIに聞いて手間を省くことで失われるのは、「誰のサイトが上手くまとまっているか？」という知識程度。それくらいなら、生成AIに聞いちゃう方が良いよね。おそらく、外国語のサイトの知識も持ってるだろうしね。
• で、AIの答えは、「要素のインデックスをハッシュテーブルで管理する」でした。
• 具体的には、以下のとおり。
• 要素をHashableとすると同時に、辞書型(ハッシュテーブル)のindicesを導入する。

/// 元のコード
struct Heap<V: Comparable>: CustomStringConvertible {
    private var data: [V] = []
    ...
}

/// 改善後コード
struct Heap<V: Comparable & Hashable>: CustomStringConvertible { // Hashableを追加
    private var data: [V] = []
    private var indices: [V: Int] = [:] // インデックスを管理するハッシュテーブルを追加
    ...
}

• 辞書型の特徴として、ハッシュテーブルを作成して要素を管理している点が上げられる。
• 配列では、indexで管理しているため、indexを指定して、要素を取り出すのは高速に行えるものの、値を指定して要素の位置を見つけるのは遅い。
• これに対して、辞書型では、キーに対してハッシュ値を算出し、ハッシュテーブルでデータ管理しているため、キーを指定すれば、すぐに要素を見つけられる。
• このindicesがどのように作られるかは、build関数に追加されたコードを見れば分かる。

/// 改善後コード
mutating func build() {
    for n in (0...(data.count / 2)).reversed() {
        heapify(n)
    }
    // インデックスを構築するためのハッシュテーブルを初期化
    for (index, value) in data.enumerated() {
        indices[value] = index // インデックスをハッシュテーブルに保存
    }
}

• 上記のとおり、「配列の値」が辞書型のキー、「配列のインデックス」が辞書型の値となる。
• 次は、追加されたremove関数がどのように動作するか確認しよう。

mutating func remove(_ value: V) -> Bool {
    // 値に対応するindexを計算量O(1)で取得、値が存在しなければ、ここで終了
    guard let index = indices[value] else {return false}

    // indexが末尾のインデックスでないとき、data[index]とdata[lastIndex]の値を入替
    let lastIndex = data.count - 1
    if index != lastIndex {
        data.swapAt(index, lastIndex) // 計算量O(1)
        indices[data[index]] = index // 入れ替え後のインデックスを更新
    }

    // 末尾のデータを削除
    data.removeLast()
    indices.removeValue(forKey: value) // ハッシュテーブルから削除

    // 並び替え実行
    if index != lastIndex {
        chgNode(index, data[index]) // 当然、chgNodeでもハッシュテーブル対応あり
    }
    
    return true
}

既存関数の変更

• heapify
  • 本来的には、indicesへの対応だけで良かったはずだけど、swapAtを使ってないことや見苦しさを指摘されたのと、そもそも、下方向にしか対応していないから、名称も変更した。なので、①swapAt導入や名称変更等と、②indicesの対応の2つに分けた。
    • それにしても、「while let」って構文は初めて知ったよ

///元のコード
mutating func heapify(_ n: Int) {
    var index = n
    var value = data[index]
    while true {
        guard let s0index = s0Index(index) else { return } // 子ノードが無ければ終了
        let s0Value = data[s0index]
        guard let s1index = s1Index(index) else { // 子ノードがs0のみの時、s0とのみ比較する
            if gt(value, s0Value) { return } // 子ノードの方が小さければ終了
            data[index] = s0Value
            data[s0index] = value
            return // s1が無いとき、下層ないためここで終了
        }
        let s1Value = data[s1index]
        let (sValue, sindex) = gt(s0Value, s1Value) ? (s0Value, s0index) : (s1Value, s1index) // どちらの子と入れ替えるか判定
        if gt(value, sValue) { return } // 子ノードの方が小さければ終了
        data[index] = sValue
        data[sindex] = value
        index = sindex // 子を親に入れ替えて、更に下の階層へ
    }
}

///①swapAt導入や名称変更等
mutating func heapifyDown(_ index: Int) { // ←名称追加(Down)、引数名をnからindexへ
    var p_index = index
    while let s0index = s0Index(p_index) { // 子ノードが無ければ終了
        let s0Value = data[s0index]
        var s_index = s0index
        if let s1index = s1Index(p_index), gt(data[s1index], s0Value) { // どちらの子と入れ替えるか判定
            s_index = s1index
        }
        if gt(data[p_index], data[s_index]) { return } // 子ノードの方が小さければ終了
        data.swapAt(p_index, s_index) // swapAt関数導入
        p_index = s_index // 子を親に入れ替えて、更に下の階層へ
    }
}

///②indicesの対応
mutating func heapifyDown(_ index: Int) {
    var p_index = index
    while let s0index = s0Index(p_index) {
        let s0Value = data[s0index]
        var s_index = s0index
        if let s1index = s1Index(p_index), gt(data[s1index], s0Value) {
            s_index = s1index
        }
        if gt(data[p_index], data[s_index]) { return }
        data.swapAt(p_index, s_index)
        indices[data[p_index]] = p_index // ←ハッシュテーブル対応
        indices[data[s_index]] = s_index // ←ハッシュテーブル対応
        p_index = s_index
    }
}

• 上方向のheapifyを追加(chgNodeで上方向に入れ替えていくため)

private mutating func heapifyUp(_ index: Int) {
    var s_index = index
    let s_value = data[s_index]
    var p_index = pIndex(s_index)

    while s_index > 0 && gt(s_value, data[p_index]) {
        data[s_index] = data[p_index]
        indices[data[p_index]] = s_index // ←ハッシュテーブル対応
        s_index = p_index
        p_index = pIndex(s_index)
    }
    data[s_index] = s_value
    indices[data[s_index]] = s_index // ←ハッシュテーブル対応
}

• あらたに、upとdownを使い分けるheapifyを作った。入れ替え後の値の大小でupとdownを切り替えるよ。

private mutating func heapify(_ index: Int, _ oldValue: V, _ newValue: V) {
    if gt(newValue, oldValue) {
        heapifyUp(index)
    } else {
        heapifyDown(index)
    }
}

• 次に、heapifyを使うchgNodeの変更に触れるね。

///元のコード
mutating func chgNode(_ i: V, _ n: Int) {
    data[n] = i
    heapify(n)
}

///変更後
mutating func chgNode(_ index: Int, _ newValue: V) { // 引数の場所を入れ替え

    // heapifyを上下方向に対応させた
    let oldValue = data[index]
    data[index] = newValue
    heapify(index, oldValue, newValue)

    // ハッシュテーブル対応
    indices[newValue] = index
    indices.removeValue(forKey: oldValue)
    
}

• insertは凄くスッキリしたよ。考えてみたら、heapifyUpって、ここで実装してたんだね。

///元のコード
mutating func insert(_ i: V) {
    data.append(i)
    var index = data.count - 1
    var value = i
    while index > 0 {
        let pindex = pIndex(index)
        let pValue = data[pindex]
        if gt(pValue, value) { return }
        data[index] = pValue
        data[pindex] = value
        index = pindex
    }
}

///変更後
mutating func insert(_ value: V) {
    data.append(value)
    indices[value] = data.count - 1 // ←ハッシュテーブル対応
    heapifyUp(data.count - 1) // heapifyUpでスッキリ！
}

• お次は、chgRootとpopRootをまとめて説明するよ

/// 元のコード
mutating func chgRoot(_ i: V) {
    chgNode(i, 0)
}

///変更後
mutating func chgRoot(_ newValue: V) {
    guard !data.isEmpty else { return } // 念のため追加(popRootにあわせた)

    // ハッシュテーブル対応
    let oldValue = data[0]
    indices[newValue] = 0
    indices.removeValue(forKey: oldValue)

    // 下方向しかないので、chgNodeから、heapifyDownに変更
    data[0] = newValue    
    heapifyDown(0)
}


///元のコード
mutating func popRoot() -> V? {
    if data.isEmpty { return nil }
    let ans = data[0]
    let last = data.removeLast()
    if data.count > 0 {
        chgRoot(last)
    }
    return ans
}

mutating func popRoot() -> V? {
    guard !data.isEmpty else { return nil }
    let rootValue = data[0]
    data[0] = data.removeLast()    
    indices.removeValue(forKey: rootValue) // ←ハッシュテーブル対応
    if !data.isEmpty {
        indices[data[0]] = 0 // ←ハッシュテーブル対応
        heapifyDown(0) // heapifyDownでスッキリ
    }
    return rootValue
}

コード全体

• 今更ながら、max値(またはmin値)を取得するコードを忘れてたのに気付いたから、rootって名前で取れることにするね。これまでは、popRootで取り去るパターンしか使ってなかったからね。
• あと、上記の関数別の比較で見落としていた部分も込みで、下記のコード全体では、修正しているよ。

struct Heap<V: Comparable & Hashable>: CustomStringConvertible {
    private var data: [V] = []
    private var indices: [V: Int] = [:]
    private var maxHeap: Bool

    var description: String { return (maxHeap ? ">" : "<") + data.description }

    var root: V? { data.isEmpty ? nil : data[0] }

    init(_ data: [V] = [], maxHeap: Bool = true) {
        self.data = data
        self.maxHeap = maxHeap
        for i in data.indices {
            indices[data[i]] = i
        }
        for i in stride(from: data.count / 2 - 1, through: 0, by: -1) {
            heapifyDown(i)
        }
    }

    private func pIndex(_ index: Int) -> Int {
        return (index + 1) / 2 - 1
    }

    private func s0Index(_ index: Int) -> Int? {
        let s0 = 2 * (index + 1) - 1
        return s0 < data.count ? s0 : nil
    }

    private func s1Index(_ index: Int) -> Int? {
        let s1 = 2 * (index + 1)
        return s1 < data.count ? s1 : nil
    }

    private func gt(_ l: V, _ r: V) -> Bool {
        return (maxHeap ? l > r : l < r)
    }

    private mutating func heapifyUp(_ index: Int) {
        var s_index = index
        let s_value = data[s_index]
        var p_index = pIndex(s_index)

        while s_index > 0 && gt(s_value, data[p_index]) {
            data[s_index] = data[p_index]
            indices[data[p_index]] = s_index
            s_index = p_index
            p_index = pIndex(s_index)
        }
        data[s_index] = s_value
        indices[data[s_index]] = s_index
    }

    private mutating func heapifyDown(_ index: Int) {
        var p_index = index
        while let s0index = s0Index(p_index) {
            let s0Value = data[s0index]
            var s_index = s0index
            if let s1index = s1Index(p_index), gt(data[s1index], s0Value) {
                s_index = s1index
            }
            if gt(data[p_index], data[s_index]) { return }
            data.swapAt(p_index, s_index)
            indices[data[p_index]] = p_index
            indices[data[s_index]] = s_index
            p_index = s_index
        }
    }

    private mutating func heapify(_ index: Int, _ oldValue: V, _ newValue: V) {
        if gt(newValue, oldValue) {
            heapifyUp(index)
        } else {
            heapifyDown(index)
        }
    }

    mutating func insert(_ value: V) {
        data.append(value)
        indices[value] = data.count - 1
        heapifyUp(data.count - 1)
    }

    mutating func popRoot() -> V? {
        guard !data.isEmpty else { return nil }
        let rootValue = data[0]
        indices.removeValue(forKey: rootValue)
        if data.count == 1 {
            return data.removeLast()
        } else {
            data[0] = data.removeLast()
            indices[data[0]] = 0
            heapifyDown(0)
            return rootValue
        }
    }

    mutating func chgNode(_ index: Int, _ newValue: V) {
        let oldValue = data[index]
        data[index] = newValue
        indices[newValue] = index
        indices.removeValue(forKey: oldValue)
        heapify(index, oldValue, newValue)
    }

    mutating func chgRoot(_ newValue: V) {
        guard !data.isEmpty else { return }
        let oldValue = data[0]
        data[0] = newValue
        indices[newValue] = 0
        indices.removeValue(forKey: oldValue)
        heapifyDown(0)
    }

    mutating func remove(_ value: V) -> Bool {
        guard let index = indices[value] else {
            return false
        }
        let lastIndex = data.count - 1
        let oldValue = data[index]
        if index != lastIndex {
            data.swapAt(index, lastIndex)
            indices[data[index]] = index
        }
        data.removeLast()
        indices.removeValue(forKey: value)
        if index != lastIndex {
            heapify(index, oldValue, data[index])
        }
        return true
    }
}

• 実行例(出力は、コメントを見て)

let vs0 = [1, 2, 3, 4, 5, 6, 7, 8, 9, 10]
let vs1 = ["c", "b", "a", "C", "B", "A", "4", "3", "2", "1"] // 降順

var h0 = Heap(vs0)
print(h0) // >[10, 9, 7, 8, 5, 6, 3, 1, 4, 2]

var h1 = Heap(vs1, maxHeap: false)
print(h1) // <["1", "2", "4", "3", "B", "A", "a", "c", "C", "b"]

// 要素の追加 (insert)
h0.insert(11)
print(h0) // >[11, 10, 7, 8, 9, 6, 3, 1, 4, 2, 5]

h1.insert("b")
print(h1) // <["1", "2", "4", "3", "B", "A", "a", "c", "C", "b", "b"]

// 要素の削除 (remove)
h0.remove(7)
print(h0) // >[11, 10, 6, 8, 9, 5, 3, 1, 4, 2]
h1.remove("4")
print(h1) // <["1", "2", "A", "3", "B", "b", "a", "c", "C", "b"]

// 根の取り出し (popRoot)
if let root0 = h0.popRoot() {
    print("Root of h0: \(root0)") // Root of h0: 11
}
print(h0) // >[10, 9, 6, 8, 2, 5, 3, 1, 4]

if let root1 = h1.popRoot() {
    print("Root of h1: \(root1)") // Root of h1: 1
}
print(h1) // <["2", "3", "A", "C", "B", "b", "a", "c", "b"]

// ノードの値変更 (chgNode)
h0.chgNode(2, 20)
print(h0) // >[20, 9, 10, 8, 2, 5, 3, 1, 4]
h1.chgNode(2, "Z")
print(h1) // <["2", "3", "Z", "C", "B", "b", "a", "c", "b"]

// 根の値変更 (chgRoot)
h0.chgRoot(15)
print(h0) // >[15, 9, 10, 8, 2, 5, 3, 1, 4]
h1.chgRoot("X")
print(h1) // <["3", "B", "Z", "C", "X", "b", "a", "c", "b"]

おわりに

• これで、ヒープを使った戦いの幅が広がったよ。

2つのheapを合体させるunion関数を作る

• ヒープは高速にmaxかminを出す目的で作るから、unionが活用される場面はないと思うけど、実験的な意味でやってみる。
• 2つのインスタンスのうち、片方のインスタンスのunionによって、新しいheapを作ることとして、2つのインスタンスに変更はないこととする。計算量は、O(n1+n2)みたいになるんだろうね。

let vs0 = [1, 2, 3, 4, 5]
let vs1 = [6, 7, 8, 9, 10]

var h0 = Heap(vs0) 
var h1 = Heap(vs1)

var hmix = h0.union(h1)

• やりたいのは、上記を可能とすること。追加したコードは、下記の通り。

    var getData : [V] {data}
    
    func union(_ other:Self) -> Self {
        Self(self.data + other.getData)
    }

• 実行結果は以下のとおり。

print(hmix) // >[10, 8, 9, 6, 7, 3, 5, 1, 4, 2]