SWEA 13041. 게으름뱅이 왕국

N명이서 K개의 작업을 수행해야 하며 각각의 사람이 수행하고자 하는 작업과 다른 작업을 하게 설득하기 위한 비용이 주어질 때 K개의 작업을 수행하기 위한 설득 비용의 합의 최솟값을 구하는 문제이다.

 

2명 이상의 사람이 참여하는 작업들에 대해서 설득 비용이 가장 큰 사람만 작업을 시키고 나머지 사람들은 설득 후보 대상이 되어 설득할 지 여부를 비용으로 정렬하여 비용이 가장 작은 사람부터 수행하지 않는 작업을 수행하여 작업의 수가 K개가 될 때 까지 설득하여 비용을 추가하면 최소가 된다. 이를 위해 우선순위 큐로 정렬하였다.

static class Persuasion implements Comparable<Persuasion> {
    int index, cost;
        
    public Persuasion(int index, int cost) {
        this.index = index;
        this.cost = cost;
    }

    @Override
    public int compareTo(Persuasion o) {
        return this.cost - o.cost;
    }
}

int[] jobs = new int[n], costs = new int[n];
st = new StringTokenizer(br.readLine());
for(int i=0; i<n; i++) jobs[i] = Integer.parseInt(st.nextToken());
st = new StringTokenizer(br.readLine());
for(int i=0; i<n; i++) costs[i] = Integer.parseInt(st.nextToken());
            
PriorityQueue<Persuasion> pq = new PriorityQueue<>();
Map<Integer, PriorityQueue<Persuasion>> map = new HashMap<>();
for(int i=0; i<n; i++) {
    PriorityQueue<Persuasion> value = map.get(jobs[i]);
    if(value == null) {
        value = new PriorityQueue<>();
        map.put(jobs[i], value);
    }
    value.add(new Persuasion(i, costs[i]));
                
    if(value.size() > 1) {
        pq.add(value.remove());
    }
}
            
int size = map.size();
long cost = 0;
while(k > size) {
    cost += pq.remove().cost;
    size++;
}
sb.append(cost).append('\n');

 

 

결과 코드는 다음과 같다.

import java.io.BufferedReader;
import java.io.IOException;
import java.io.InputStreamReader;
import java.util.HashMap;
import java.util.Map;
import java.util.PriorityQueue;
import java.util.StringTokenizer;

public class Solution {

    public static void main(String[] args) throws IOException {
        BufferedReader br = new BufferedReader(new InputStreamReader(System.in));
        int T = Integer.parseInt(br.readLine());
        
        StringBuilder sb = new StringBuilder();
        for(int tc=1; tc<=T; tc++) {
            sb.append('#').append(tc).append(' ');
            StringTokenizer st = new StringTokenizer(br.readLine());
            int n = Integer.parseInt(st.nextToken()), k = Integer.parseInt(st.nextToken());
            int[] jobs = new int[n], costs = new int[n];
            st = new StringTokenizer(br.readLine());
            for(int i=0; i<n; i++) jobs[i] = Integer.parseInt(st.nextToken());
            st = new StringTokenizer(br.readLine());
            for(int i=0; i<n; i++) costs[i] = Integer.parseInt(st.nextToken());
            
            PriorityQueue<Persuasion> pq = new PriorityQueue<>();
            Map<Integer, PriorityQueue<Persuasion>> map = new HashMap<>();
            for(int i=0; i<n; i++) {
                PriorityQueue<Persuasion> value = map.get(jobs[i]);
                if(value == null) {
                    value = new PriorityQueue<>();
                    map.put(jobs[i], value);
                }
                value.add(new Persuasion(i, costs[i]));
                
                if(value.size() > 1) {
                    pq.add(value.remove());
                }
            }
            
            int size = map.size();
            long cost = 0;
            while(k > size) {
                cost += pq.remove().cost;
                size++;
            }
            sb.append(cost).append('\n');
        }
        
        br.close();
        System.out.print(sb);
    }

    static class Persuasion implements Comparable<Persuasion> {
        int index, cost;
        
        public Persuasion(int index, int cost) {
            this.index = index;
            this.cost = cost;
        }

        @Override
        public int compareTo(Persuasion o) {
            return this.cost - o.cost;
        }
    }
}