【AC5 解説参考に自力実装】C - Variety

振り返り

  • 答えそのものは公式の解説しか見ずに理解しているコードでACできたので嬉しい
  • 細切れで5~6日くらいかけた
  • 生成AIには貪欲法の概念を聞いただけだった
  • 地道にプリントデバッグしたらおかしいところも見えた
  • for i in range(len(arrayname)): みたいなのに頼らず、論理的をちゃんと設計してMとかを使ってACしたい

コード

# 標準入力の1行目を取得
N, K, M = map(int, input().split())

# 解説を参考
## 個数分+1この[]を配列に入れる。indexを色、価値を配列内の値として登録
jewels = [[] for i in range(N + 1)]

## 個数分ループを回して、2行目以降の標準入力から、jewelsへ適切な色のindexで、価値を数値で配列内に登録
## [[], [30, 40, 50], [10], [20], [], []] のようになる
for i in range(N):
    C, V = map(int, input().split())
    jewels[C].append(V)
  
colorfv = []  
# 取らなければいけない色数分、色の配列から先にとる準備のために、各色の先頭だけの配列に詰める
# 準備するのにcolorfvに抜き出したものは、元の宝石の配列内で印をつける
for i in range(len(jewels)):
  if (jewels[i]):
    jewels[i].sort(reverse=True)
    colorfv.append(jewels[i][0])
    jewels[i][0]='x'
colorfv.sort(reverse=True)

# 元の宝石の配列を平坦化し、印をつけたものを削除し、大きい順にソート
flatjewels = []
for i in range(len(jewels)):
  for j in range(len(jewels[i])):
    flatjewels.append(jewels[i][j])
flatjewels = [i for i in flatjewels if i != 'x']
flatjewels.sort(reverse=True)

# 色のノルマ達成する
result = []
for i in range(len(colorfv)):
  result.append(colorfv[i])
  colorfv[i]='x'
  if (M<=len(result)): ## ノルマ達成したら抜ける
    break

# 色ノルマ用のデータと他全部結合
flatjewels = [i for i in (flatjewels + colorfv) if i != 'x']
flatjewels.sort(reverse=True)

# 色のノルマ達成してもKにたりない個数は、大きいものから順番に追加
for i in range(K-len(result)): # 論理的矛盾がなければ、K-Mとかでいいはず?
  if (K<=len(result)):
    break
  result.append(flatjewels[i])

print(sum(result))