【TLE2 解説読み実装直し】C - Variety

やったこと

  • 貪欲法について
    • 問題の解説からgeminiにおすすめしてもらった
    • notebooklmで簡単に概念を聞いた

振り返り

  • 貪欲法の理解が浅いので純粋なインプットをしてから再挑戦する

コード

  • 相変わらずTLEがあるのと、WAが1件ある。
N, K, M = list(map(int, input().split()))
l = [list(map(int, input().split())) for l in range(N)]
result=0
clrv=[]
li = sorted(l, reverse=True, key=lambda x: x[1])
k=[]
v=[]

for i in range(N):
  if (len(set(k))>=M):
    print(sum(v))
    break
  for i in range(K):
    if (len(li)>0):
      v.append(li[i][1])
      k.append(li[i][0])
    if K-len(k) == 0:
      break
  if (len(set(k)) < M):
    for i in range(M-len(set(k))):
      del v[-1]
      del k[-1]
  if (len(set(k))>=M):
    print(sum(v))
    break
  for i in range(N):
    if (len(li) and li[0][0] in k):
      del li[0]

メモ

N, K, M = list(map(int, input().split()))
l = [list(map(int, input().split())) for l in range(N)]
result=0
clrv=[]
li = sorted(l, reverse=True, key=lambda x: x[1])
k=[]
v=[]

for i in range(N):
  # print("")
  # print(f"=={i+1}周目==")
  # print(f"li:{li}")
  # print(f"v:{v}")
  # print(f"len(set(k)):{len(set(k))}")
  # チェックして条件満たしていれば出力
  if (len(set(k))>=M):
    print(sum(v))
    break
  for i in range(K):
    if (len(li)>0):
      v.append(li[i][1])
      k.append(li[i][0])
    if K-len(k) == 0:
      break
  # 色がM種類なかったら足りない種類の数、選んだものを価値が小さい順に捨てる
  if (len(set(k)) < M):
    for i in range(M-len(set(k))):
      del v[-1]
      del k[-1]
      # print(f"{M-len(set(k))}色たりない")
      # print(k)
      # print(v)
  # チェックして条件満たしていれば出力
  if (len(set(k))>=M):
    print(sum(v))
    break
  # sortedした取得元のリストから、とっている色のものを削除
  for i in range(N): # 取る宝石の数走査できれば十分
    # print(f"out of rangeするifの手前li:{li}")
    # print(f"消す走査前のli:{li}")
    # print(f"消す走査前のk:{k}")
    # print(f"消す走査前のv:{v}")
    if (len(li) and li[0][0] in k): # delでindexも詰まるので[0][0]でliの先頭の価値をとる
      # print(i)
      # print(f"{li[0]}を消す対象として見ている")
      del li[0]
      # print(f"消した後li:{li}")
  # print(k)