ABC 116 D - Various Sushi (Go)
irisruneです。回転寿司なんかに行くとできるだけ多くの種類のネタを食べる派です。
多分実装が難しい部類の問題じゃないでしょうか。結構ゴリゴリ書いたので長めです。
package main
import (
"fmt"
"sort"
)
type susi struct {
t int
d int
}
type sushi []susi
func (s sushi) Len() int {
return len(s)
}
func (s sushi) Swap(i, j int) {
s[i], s[j] = s[j], s[i]
}
func (s sushi) Less(i, j int) bool {
return s[i].d < s[j].d
}
func maxInt(a, b int) int {
switch {
case a > b:
return a
default:
return b
}
}
func main() {
var n, k int
fmt.Scan(&n, &k)
var sushi sushi = make([]susi, n)
for i := range sushi {
fmt.Scan(&sushi[i].t, &sushi[i].d)
}
sort.Sort(sort.Reverse(sushi))
pointD := 0
typeNum := 0
types := make([]int, n+1)
for i := 0; i < k; i++ {
pointD += sushi[i].d
types[sushi[i].t]++
if types[sushi[i].t] == 1 {
typeNum++
}
}
pointSum := pointD + (typeNum * typeNum)
iRemove := k - 1
iAdd := k
for typeNum < k {
for ; iRemove > 0; iRemove-- {
if types[sushi[iRemove].t] > 1 {
pointD -= sushi[iRemove].d
types[sushi[iRemove].t]--
break
}
}
if iRemove == 0 {
break
}
iRemove--
for ; iAdd < n; iAdd++ {
if types[sushi[iAdd].t] == 0 {
pointD += sushi[iAdd].d
types[sushi[iAdd].t]++
typeNum++
break
}
}
if iAdd == n {
break
}
iAdd++
pointSum = maxInt(pointSum, pointD+(typeNum*typeNum))
}
fmt.Println(pointSum)
}
今回は構造体の配列とそのソートまで実装して解いてみました。配列を作ることすらままなかった頃に比べると色々できるようになってきたと思います。
アルゴリズムの大筋は解説の通りなのですが、このコードの特徴は優先度付きキューなどを用いていないところです。 そのためどのようにして解いているかというと、
についてネタ
を食べた数を記録しておく。
- おいしさの総和とネタの種類数を求めておく。
個の寿司をおいしさで降順ソートしておき、取り除く寿司を探索するインデックスを記録しておく。
- さらに、追加する寿司を探索するインデックスも記録しておく。
- 記録したインデックスを用いて食べる寿司を入れ替える操作を行い、おいしさの総和とネタの種類数を更新する。
ポイントは、食べた寿司そのものを記録するのではなく、おいしさの総和とそれぞれのネタを食べた数を記録しておくことです。
これにより、全体の計算量は優先度付きキューを用いた場合よりさらに少ないになる…はずです、多分。
コードレビュー
- ほとんどの処理を関数に退避していないのでmain関数が肥大化している。
- ループ部分と独立して最初においしい寿司だけを食べる部分もあるのでさらに肥大化している。
- 計算量削減のための変数iRemove,iAddの扱いが特殊でわかりにくい。
sushiに複数形がないから構造体susiの配列としてsushiを定義するというよくわからないことになっている。
雑記
- 1114msかかってるのは多分入力のせいだと思うのでそろそろfmt.Scan以外も使えるようにしておきたいですね。
