AtCoder:ABC116-D(Various Sushi)で苦戦したので復習した

こんにちは、ぐぐりら(<a href="https://twitter.com/guglilac">@guglilac</a>)です。

何気に初めての復習記事かもしれません。
AtCoderのABC116に参加したのですが、Cまでしか解けなかった(しかもBで焦ってWA)
ので、Dを復習してみました。

<a href="https://atcoder.jp/contests/abc116/tasks/abc116_d">D - Various Sushi</a>

### コンテスト中
種類数がめんどそうだから種類の数を固定してそれぞれ実験すればいいのかな、と考えました。

それはあってたぽいんですが、そこからがわかりませんでした。

種類の個数を固定した時に得られる最大美味しさポイントをそれぞれ求めればいいわけですが、

どの種類のネタを使うかも調べてたら組み合わせ多すぎて計算終わらんやん、とか、アホなこと考えてました。

1種類ずつ増やしていく、というのが正解なわけですが、コンテスト中に思いついたのは本当に愚直に1種類から増やしていくもので、これで考えていたのでswapする発想にたどり着けませんでした。。。

いや、1種類から増やしていっても、想定解法の貪欲にとっていったときの種類数から増やしていってももちろん同じですが、1種類だからわかりづらかったという言い訳です。

そもそも必ずしも1種類になるように選べるわけじゃないので、これじゃない感を感じてしまいました。。

### 解説見てから

解説みてみて、種類ごとに決め打ちしてやっていくのはあってたのかーとなりました。

1種類ずつ増やしていくやり方は、解説にもある通り

* 現時点で食べることになっている寿司から、食べなくても種類数が減らない寿司のうち美味しさ最小の寿司

と

* 現時点で食べることになっていない寿司から、新しい種類の寿司のうち美味しさ最大の寿司

をswapしていくと、種類数を一つ増やしたときの最大美味しさポイントが得られます。

理屈はわかったので早速コンテスト後に実装してみようと思い、解説にあったようにpriority queueを使って実装してみようと思ったのですが、うまくできませんでした。

がーっと書いてみて、swapするところで詰まってしまいました。

swapする候補の寿司集合をpriority queueで表現して、取り出したり挿入したりすればいいのかと思いましたが、
そのpriority queueが各種類数のときの処理ごとに作成し直す必要があるように感じました。

下のプログラム中のQとRがそれぞれ
* 現時点で食べることになっている寿司のうち、食べなくても種類数が減らない寿司

と

* 現時点で食べることになっていない寿司のうち、新しい種類の寿司

を表しています。
ある寿司のペアをswapした時に、QとRはその寿司だけpopとpushをすればいいわけではなく、毎回種類ごとに寿司がいくつあるかを見てQとRを構成しないといけないような気がしてしまいます。

そうすると、種類の個数は$O(N)$だから全体で$O(N^2)$(種類数でイテレート)かかってアウトなんじゃないかという疑問です。

priority queueを使う実装は結局通せなかったので、強い方に教えていただきたいです。。。

```python
from collections import Counter
from heapq import heappush, heappop
N, K = map(int, input().split())
sushi_list = []
for i in range(N):
    t, d = map(int, input().split())
    sushi_list.append((t, d))

sushi_list.sort(key=lambda x: x[1], reverse=True)
S_x = sushi_list[:K]
c = Counter()
f_p = 0

for t, d in S_x:
    c[t] += 1
    f_p += d

Q = []
R = []

for sushi in sushi_list:
    if sushi[0] in c.keys():
        if c[sushi[0]] > 1:
            heappush(Q, sushi[1])
    else:
        heappush(R, -sushi[1])
p = len(c)
ans = f_p + p**2

while True:
    print(Q, R)
    if len(Q) == 0 or len(R) == 0:
        break
    min_q = heappop(Q)
    max_r = -heappop(R)
    p += 1
    f_p = f_p + max_r - min_q
    ans = max(ans, f_p + p**2)

print(ans)
```

### 通った！

結局、やり方を変えました。
priority queueは使わず、sortしたsushi_listとは別に、それに対応したgood配列というものを定義しました。

今食べることになっている寿司のうち、swapの対象にするかしないかを表す配列です。

good配列は、sushi_iがそのネタの中で最大の美味しさかどうかを表すbool配列です。
ここの前処理は$O(N)$でできます。

下のプログラムでいう`max_d[type]`という、 各ネタの中で最大の美味しさポイントを格納する配列を用意して、それとの比較をしていくとできます。

この、

*あるかわからないネタの分も配列を用意する*

というのが未だにしっくりこなくてすぐに思いつかないのが反省点です。

速くしたり簡単に書こうと思ったら、どかっと用意するのも必要という学びでした。
(ネタでグループに分けてその中で最大値とネタをとってきて辞書にして、、とか考えて難しく書こうとしてしまいました)

ここで、注意点が一つあって、

good配列を作る際に、同じネタで最高の美味しさポイントが複数ある場合があります。
普通にmax_d配列と比較して最大ならgood配列に`True`を入れると、そのような寿司全部がgoodになってしまいます。

そうすると、同じネタで複数の最高美味しさポイントの寿司があったときに、swapの対象になってくれません。

各ネタの寿司でgoodな寿司はただ1つであって欲しいので、一番最初に出会った最大おいしさポイントの寿司のみgoodにします。

下のプログラムでいうとコメントアウトでここ！！！と書いてある部分で、そこを修正してます。
一度goodにしたsushiのネタにはもうマッチしないように適当にありえない数値を入れておきます。

これを忘れて1WAだしてしまいました。

かなしい。

あとはsortしたsushi_listについて、swapする対象を探すインデックスrとlを減らし/増やして行ってgoodかどうか見に行ってswap、そのときのポイントの合計を比較してmaxをとる、みたいな感じでできます。

swapした後はちゃんとインデックスを一つ進めるのを忘れずに。(whileから抜けてくれません)

```python
N, K = map(int, input().split())
sushi_list = []
max_d = [0] * N
for i in range(N):
    t, d = map(int, input().split())
    sushi_list.append((t - 1, d))
    max_d[t - 1] = max(max_d[t - 1], d)

sushi_list.sort(key=lambda x: x[1], reverse=True)
good = []
for i in range(N):
    if max_d[sushi_list[i][0]] == sushi_list[i][1]:
        good.append(True)
        max_d[sushi_list[i][0]] = 0 # ここ！！！
    else:
        good.append(False)

s = sum([sushi[1] for sushi in sushi_list[:K]])
p = len(set([sushi[0] for sushi in sushi_list[:K]]))
ans = s + p**2
r = K - 1
l = K
while True:
    while r >= 0 and good[r]:
        r -= 1
    while l < N and not good[l]:
        l += 1
    if r < 0 or l >= N:
        break
    s = s - sushi_list[r][1] + sushi_list[l][1]
    p += 1
    ans = max(ans, s + p**2)
    r -= 1
    l += 1

print(ans)
```

### おわりに

書き直してACしたのですっきりはしたのですが、優先度つきキューでどうやって実装するのかいまいちピンときていません。

実装できをつける部分が多くて大変な問題だなあ、としんどい気持ちになりました。これで400点なんですね。

またしっかり復習した問題があればブログに記事を書きたいと思います。

このブログを検索

ぐぐりらにっき

AtCoder:ABC116-D(Various Sushi)で苦戦したので復習した

コメント

コメントを投稿

AtCoder:ABC116-D(Various Sushi)で苦戦したので復習した

この記事をシェアする

コメント

コメントを投稿