AtCoder ABC141 D - Powerful Discount Tickets (Go) - あるエンジニアのAtCoder奮闘記

irisruneです。D問題が大体安定して解けるようになっている辺り成長を感じますね。

問題

Go言語にとってソート以上の鬼門と言えるであろうヒープを使う問題です。

package main

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

var sc *bufio.Scanner

func nextStr() string {
    sc.Scan()
    return sc.Text()
}

func nextInt() int {
    sc.Scan()
    i, e := strconv.Atoi(sc.Text())
    if e != nil {
        panic(e)
    }
    return i
}

type Heap []int

func (h Heap) Len() int {
    return len(h)
}

func (h Heap) Less(i, j int) bool {
    return h[i] > h[j] // reversed-Heap
}

func (h Heap) Swap(i, j int) {
    h[i], h[j] = h[j], h[i]
}

func (h *Heap) Push(x interface{}) {
    *h = append(*h, x.(int))
}

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

func main() {
    sc = bufio.NewScanner(os.Stdin)
    sc.Buffer(make([]byte, 0), 1000000001*3)
    sc.Split(bufio.ScanWords)
    n, m := nextInt(), nextInt()

    h := &Heap{}
    heap.Init(h)
    for i := 0; i < n; i++ {
        heap.Push(h, nextInt())
    }

    for j := 0; j < m; j++ {
        (*h)[0] /= 2
        heap.Fix(h, 0)
    }

    ans := 0
    for _, v := range *h {
        ans += v
    }
    fmt.Println(ans)
}

ある商品に対して $Y$ 枚の割引券を使う操作は、 $1$ 枚の割引券を使った後に $Y-1$ 枚の割引券を使う操作と等価です（厳密な証明は公式解説を参照してください）。これにより割引券を1枚ずつ、その時点で最も高い商品に対して使う手法で解を求めることができます。

その際必要な計算量について考えると、愚直に実装した場合は最も高い商品を探す処理に $O(N)$ かかるため、全体で $O(NM)$ かかってしまい間に合いません。つまり最も高い商品を探す処理に $O(\log N)$ 以下の計算量となるように実装する必要があると考えられます（一応 $O(\sqrt{N})$ でもよいですが、探索・ソートアルゴリズムの計算量ではあまり見かけないと思います）。

ここでヒープ（あるいは優先度付きキュー）を用いることで全体の計算量を $O((N+M)\log N)$ に抑えることができます。これはヒープの初期化処理が $O(N\log N)$ 、要素の追加・削除（更新）が $O(\log N)$ 、最大要素の取得が $O(1)$ の計算量で行えるためです。

ところで、他の方法として仮に二分探索を用いるとすると、以下のような実装が考えられます。

全商品をソートする。
最も高い商品に割引券を使う。
割引券を使った商品が入る場所を二分探索によって求める。
求めた場所に割引券を使った商品を入れる。
割引券がある限り2.～4.を繰り返す。

この実装は一見問題ないように思えますが、4.の配列に対する挿入処理で $O(N)$ の計算量がかかってしまいます。そのため結局全体の計算量が $O(NM)$ となってしまうのでこの実装は使えません。

結局、配列を用いる場合は探索あるいは挿入がボトルネックになってしまうので、配列以外のデータ構造を用いることになると思います。

雑記

Goでヒープを使うためには、単純にソートインタフェース以上の数のメソッドを定義する必要があり手間がかかってしまいます。
- ソートは最近のバージョンで実装の簡略化が可能になりましたが、ヒープは特にそういうこともないので使わずに済むならそれに越したことはないですね。
今週はABC141-Eの嘘解法か最近解けた1300点問題かのどちらかについて書く予定です。~~前者の方が書きにくそうなので折れたら後者を書きます。~~