[백준 23889번] 돌 굴러가유

23889번: 돌 굴러가유

from itertools import combinations

n,m,k = map(int,input().split())

# i번째 마을의 모래성의 개수 xi
house_count = [0]+list(map(int,input().split()))

# 돌이 굴러가기 시작하는 마을의 위치
rock_locations = list(map(int,input().split()))

house_status = [True]*(n+1)

# 벽의 개수 m개의 줄에 걸쳐 가장 많은 모래성을 지키기 위하여 
# 벽을 설치해야 할 마을의 위치를 오름차순으로 출력
# (2가지 이상의 경우에는 사전순으로 가장 빠른 답을 출력)

for rock in rock_locations:
    house_status[rock] = False

cases = []

for start in rock_locations:
    total = house_count[start]
    for i in range(start+1,n+1):
        if not house_status[i]:
            break
        total += house_count[i]
    cases.append((total,start))

cases.sort(key=lambda x:(x[0],-x[1]))

answer = sorted([x[1] for x in cases[-m:]])

for i in answer:
    print(i)

이번 문제는 각 마을마다 존재하는 모래성 개수와 돌이 굴러가기 시작하는 마을 인덱스가 주어졌을 때, 최대한 많은 모래성을 지키기 위해서 선택해야 할 마을의 위치를 오름차순으로 출력하는 문제이다.

 

처음에 이 문제를 풀 때 나도 모르게 Greedy하게 풀기 위하여 노력했었는데, 문제를 풀이하는 방식이 다소 복잡해서 바로 풀이하지 못하였다.

 

그냥 단순하게 각 돌이 굴러가기 시작하는 마을 인덱스부터 모래성 개수를 세어 지킬 수 있는 모래성의 개수의 총합을 구하고, 그 총합을 기준으로 정렬한 뒤 뒤에서 m개를 자르면 된다.

 

하지만 여기서 조심해야할 점은 "가장 많은 모래성을 지킬 수 있는 경우가 두 가지 이상 존재할 경우, 사전순으로 가장 빠른 답을 출력한다."라는 조건이다.

 

위 조건 때문에 정렬해줄 때 지킬 수 있는 모래성의 개수가 같다면 인덱스를 작은 것을 선택해야 하기 때문에 오름차순으로 정렬해주어야 한다.

 

최대나 최소를 구하는 문제가 나왔을 때 그리디한 방법을 떠올리고, 간단하게 풀이할 수 있는 방법들을 떠올리는 연습을 해야할 것 같다.

 

배운 점 및 느낀 점

1. 최대나 최소를 구할 때는 완전탐색 말고도 그리디한 방법으로도 풀 수 있다.
2. 두가지 이상의 조건으로 정렬하기 위해서는 lambda x:(x[0],-x[1]) .. 이런 방식으로 정렬할 수 있다.