[알고리즘, 코드트리] 코드트리 투어 - java

문제

문제가 상당히 길다. 전문을 보고 싶은 사람은 코드트리에서 확인하자.

코드트리 | 코딩테스트 준비를 위한 알고리즘 정석

요약하자면 이렇다. 총 다섯 가지 동작을 수행해야 한다.

(1) 코드트리 랜드 관련 정보가 주어집니다. 도시의 수 n 과 간선의 수 m, 그리고 m 개의 간선에 해당하는 정보 (vi​, ui​, wi​) 가 주어집니다. 

(2) 코드트리 여행사는 (id, revenueid​, destid​) 에 해당하는 여행 상품을 추가로 만들고 이를 관리 목록에 추가

(3) 고유 식별자 id 에 해당하는 여행 상품이 존재하는 경우, 해당 id 의 여행 상품을 관리 목록에서 삭제

(4) 관리 목록에서 조건에 맞는 최적의 여행 상품을 선택하여 판매, 만약 판매 가능한 상품이 전혀 없다면 −1 을 출력하고 상품을 제거하지 않게 됩니다.

(5) 여행 상품의 출발지를 전부 s 로 변경

입력 예시

출력 예시

문제 풀이

이 문제 역시 지문을 천천히 잘 읽으면 어렵지 않게 구현해 낼 수 있다.

자료 구조 정의 및 입출력

static class Edge implements Comparable<Edge> {
    int u;
    int w;

    Edge(int u, int w) {
        this.u = u;
        this.w = w;
    }

    @Override
    public int compareTo(Edge o) {
        return w - o.w;
    }

    public String toString() {
        return "연결된 노드는 " + u + " || 가중치는 " + w;
    }
}

static class Item implements Comparable<Item> {
    int id;
    int revenue;
    int dest;
    int cost;

    Item(int id, int revenue, int dest) {
        this.id = id;
        this.revenue = revenue;
        this.dest = dest;
        this.cost = revenue - costArr[dest];
    }

    @Override
    public int compareTo(Item o) {
        if(o.cost < 0) {
            return -1;
        } else if(cost < 0) {
            return 1;
        } else if(o.cost - cost == 0) {
            return id - o.id;
        }

        return o.cost - cost;
    }
}

static ArrayList<ArrayList<Edge>> g = new ArrayList<>();
static Map<Integer, Item> itemMap = new HashMap<>();
static PriorityQueue<Item> itemPq = new PriorityQueue<>();
static int[] costArr;

선언할 자료구조가 많다. 우선 그래프 형태를 표현하기 위해서 Edge class를 작성했다('Node'와 의미상 다를게 크게 없다). 또한, 여행 상품을 관리하기 위한 Item class도 작성해 두었다. 추후에 있을 연산때문에 Edge와 Item에는 모두 Comparable<>을 붙여놓았다. 이유는 후술한다.

그래프의 간선 정보를 저장하기 위해서 g를 선언하고, 여행상품을 관리하기 위해서 itemMap과 itemPq를 동시에 선언한다. 이 역시 이유는 후술하겠다.

int Q = Integer.parseInt(br.readLine());
int v, u, w;
int id, revenue, dest, s;
while(Q-- > 0) {
    st = new StringTokenizer(br.readLine());
    int command = Integer.parseInt(st.nextToken());

    switch(command) {
        case 100 :
            // 코드트리 랜드 건설, 항상 첫번째로 1회만 실행
            // makeLand();
            break;
        case 200 :
            // 여행 상품 생성
            // createItem();
            break;
        case 300 :
            // 여행 상품 취소
            // cancelItem(id);
            break;
        case 400 :
            // 최적 여행 상품 판매
            // sell();
            break;
        case 500:
            // 출발지 변경
            // changeStart();
            break;
    }
}

문제에서 정의된 대로 흐름을 분기하여 구조를 짰으므로 각각 필요한 함수를 작성해보자.

command 100 : 코드트리 랜드 건설 _ makeLand()

static private void makeLand(int n, int m, StringTokenizer st) {
    for(int i = 0; i < n; i++) {
        g.add(new ArrayList<Edge>());
    }

    for(int i = 0; i < m; i++) {
        int u = Integer.parseInt(st.nextToken());
        int v = Integer.parseInt(st.nextToken());
        int w = Integer.parseInt(st.nextToken());

        g.get(u).add(new Edge(v, w));
        g.get(v).add(new Edge(u, w));
    }

    changeStart(0);
}

사실상 초기화 함수의 역할을 하는 makeLand 함수이다. 테스트 케이스에서 가장 처음 명령으로 단 1회 주어지게 되므로 이 함수에서 graph 간선 정보를 저장하였다. 마지막으로 최초 출발지(0)에서 모든 노드 사이의 최소 거리를 구해야 하므로 500 명령으로 수행할 changeStart() 함수를 실행한다.

command 500 : 출발지 변경 _ changeStart()

static private void changeStart(int s) {
    // Dijkstra로 모든 여행 상품까지 COST 계산
    // 다익스트라 알고리즘 초기화
    costArr = new int[n]; // 최소 비용을 저장할 배열
    for (int i = 0; i < n; i++) {
        costArr[i] = Integer.MAX_VALUE;
    }

    PriorityQueue<Edge> pq = new PriorityQueue<Edge>();
    pq.offer(new Edge(s, 0));
    costArr[s] = 0;
    while (!pq.isEmpty()) {
        Edge now = pq.poll();

        if (costArr[now.u] < now.w) {
            continue;
        }

        // 선택된 노드의 모든 주변 노드를 고려한다.
        for (int i = 0; i < g.get(now.u).size(); i++) {
            Edge next = g.get(now.u).get(i);
            if (costArr[next.u] > now.w + next.w) {
                costArr[next.u] = now.w + next.w;
                pq.add(new Edge(next.u, costArr[next.u]));
            }
        }
    }

    PriorityQueue<Item> temp = new PriorityQueue<>();
    while(!itemPq.isEmpty()) {
        Item item = itemPq.poll();
        Item newItem = new Item(item.id, item.revenue, item.dest);
        temp.add(newItem);
    }
    itemPq = temp;
}

파라미터로 주어진 s에서 시작하여 모든 노드 사이에 최소 거리를 구하기 위해서 Dijkstra 알고리즘을 사용한다. 이 때, 시간복잡도를 최소화하기 위해서 가중치가 가장 작은 Edge를 먼저 선택하는 heap이 필요하다(자료구조에서 Edge에 Comparable을 구현하고 있는 이유이다).

또한 최단 거리의 기준점이 바뀌어서 여행 상품의 최적을 구하는 값도 변경되었기 때문에 이를 반영하기 위해서 cost를 재계산(Item 객체 생성 시점에 생성자에 의해 재계산)하고 새로운 pq에 값을 옮겨 담아준다.

command 200 : 여행 상품 생성 _ createItem()

static private void createItem(int id, int revenue, int dest) {
    itemMap.put(id, new Item(id, revenue, dest));
    itemPq.add(itemMap.get(id));
}

여행 상품을 생성하고 itemMap과 itemPq에 각각 따로 저장한다.

command 300 : 여행 상품 취소 _ cancelItem()

static private void cancelItem(int id) {
    if(itemMap.get(id) == null) {
        return;
    }
    itemMap.remove(id);
}

문제를 읽어보면 다음과 같이 설명한다. "고유 식별자 id 에 해당하는 여행 상품이 존재하는 경우 해당 id 의 여행 상품을 관리 목록에서 삭제해야 함을 의미합니다." 뉘앙스가 존재하지 않는 id를 입력하는 상황도 가정한 느낌이라 조건문으로 itemMap에서 조회에 실패하면 함수를 종료하도록 설정했다.

여행 상품 취소의 경우 itemMap에 대해서만 Item을 삭제한다. itemPq에서의 삭제 상품 처리는 command 400에서 후술한다.

command 400 : 최적의 여행상품 판매 _ sell()

static private void sell() {
    if(itemPq.isEmpty()) {
        System.out.println(-1);
        return;
    }

    Item item = itemPq.poll();
    while(!itemPq.isEmpty() && itemMap.get(item.id) == null) {
        item = itemPq.poll();
    }
    
    if(itemMap.get(item.id) == null) {
        System.out.println(-1);
        return;
    }


    if(item.cost < 0) {
        System.out.println(-1);
        itemPq.add(item);
    }else {
        itemMap.remove(item.id);
        System.out.println(item.id);
    }
}

itemPq에서 Item을 하나 꺼내고 itemMap에서 삭제한다. 단 조건 분기가 까다롭다.

(1) itemPq의 상태가 비어있다면 판매할 상품이 없다.

(2) itemPq에서 Item을 꺼내어도 itemMap에 없다면 이미 삭제된 상품이다. 차순위 Item을 꺼내자.

(3) itemPq에서 최적의 Item을 꺼내도 기준값(revenue - cost)이 음수라면 판매 불가 상품이다. 도로 itemPq에 넣어 놓고 판매할 상품이 없음을 표시하자.

 

이 연산을 최적화하기 위해서 Item의 compareTo() 비교 조건을 아래와 같이 주었다.

static class Item implements Comparable<Item> {
    int id;
    int revenue;
    int dest;
    int cost;

    Item(int id, int revenue, int dest) {
        this.id = id;
        this.revenue = revenue;
        this.dest = dest;
        this.cost = revenue - costArr[dest];
    }

    @Override
    public int compareTo(Item o) {
        if(o.cost < 0) {
            return -1;
        } else if(cost < 0) {
            return 1;
        } else if(o.cost - cost == 0) {
            return id - o.id;
        }

        return o.cost - cost;
    }
}

(1) 어떤 경우에도 음수는 가장 낮은 우선순위를 가져야 하므로 비교 객체와 대조 객체의 기준값(cost)이 음수이면 -1을 반환다.

(2) cost가 동일하면 id가 작은 Item이 우선순위를 가진다.

(3) cost는 클수록 높은 우선순위를 가진다.

전체 코드

import java.util.*;
import java.io.*;

public class Main {

    static class Edge implements Comparable<Edge> {
        int u;
        int w;

        Edge(int u, int w) {
            this.u = u;
            this.w = w;
        }

        @Override
        public int compareTo(Edge o) {
            return w - o.w;
        }

        public String toString() {
            return "연결된 노드는 " + u + " || 가중치는 " + w;
        }
    }

    static class Item implements Comparable<Item> {
        int id;
        int revenue;
        int dest;
        int cost;

        Item(int id, int revenue, int dest) {
            this.id = id;
            this.revenue = revenue;
            this.dest = dest;
            this.cost = revenue - costArr[dest];
        }

        @Override
        public int compareTo(Item o) {
            if(o.cost < 0) {
                return -1;
            } else if(cost < 0) {
                return 1;
            } else if(o.cost - cost == 0) {
                return id - o.id;
            }

            return o.cost - cost;
        }
    }

    static class Cost {
        int cost;

        Cost(int cost) {
            this.cost = cost;
        }

        int getCost() {
            return cost;
        }

        void setCost(int cost) {
            this.cost = cost;
        }
    }

    static ArrayList<ArrayList<Edge>> g = new ArrayList<>();
    static Map<Integer, Item> itemMap = new HashMap<>();
    static PriorityQueue<Item> itemPq = new PriorityQueue<>();
    // 객체 참조 할 것
    static int[] costArr;

    static int n, m;

    public static void main(String[] args) throws IOException {
        BufferedReader br = new BufferedReader(new InputStreamReader(System.in));
        StringTokenizer st;

        int Q = Integer.parseInt(br.readLine());
        int v, u, w;
        int id, revenue, dest, s;
        while(Q-- > 0) {
            st = new StringTokenizer(br.readLine());
            int command = Integer.parseInt(st.nextToken());

            switch(command) {
                case 100 :
                    n = Integer.parseInt(st.nextToken());
                    m = Integer.parseInt(st.nextToken());

                    makeLand(n, m, st);
                    break;
                case 200 :
                    id = Integer.parseInt(st.nextToken());
                    revenue = Integer.parseInt(st.nextToken());
                    dest = Integer.parseInt(st.nextToken());

                    createItem(id, revenue, dest);
                    break;
                case 300 :
                    id = Integer.parseInt(st.nextToken());

                    cancelItem(id);
                    break;
                case 400 :
                    sell();
                    break;
                case 500:
                    id = Integer.parseInt(st.nextToken());

                    changeStart(id);
                    break;
            }
        }
    }

    static private void makeLand(int n, int m, StringTokenizer st) {
        for(int i = 0; i < n; i++) {
            g.add(new ArrayList<Edge>());
        }

        for(int i = 0; i < m; i++) {
            int u = Integer.parseInt(st.nextToken());
            int v = Integer.parseInt(st.nextToken());
            int w = Integer.parseInt(st.nextToken());

            g.get(u).add(new Edge(v, w));
            g.get(v).add(new Edge(u, w));
        }

        changeStart(0);
    }

    static private void createItem(int id, int revenue, int dest) {
        itemMap.put(id, new Item(id, revenue, dest));
        itemPq.add(itemMap.get(id));
    }

    static private void changeStart(int s) {
        // Dijkstra로 모든 여행 상품까지 COST 계산
        // 다익스트라 알고리즘 초기화
		costArr = new int[n]; // 최소 비용을 저장할 배열
		for (int i = 0; i < n; i++) {
			costArr[i] = Integer.MAX_VALUE;
		}

        PriorityQueue<Edge> pq = new PriorityQueue<Edge>();
		pq.offer(new Edge(s, 0));
		costArr[s] = 0;
		while (!pq.isEmpty()) {
			Edge now = pq.poll();

			if (costArr[now.u] < now.w) {
				continue;
			}

			for (int i = 0; i < g.get(now.u).size(); i++) {
				Edge next = g.get(now.u).get(i);
				if (costArr[next.u] > now.w + next.w) {
					costArr[next.u] = now.w + next.w;
					pq.add(new Edge(next.u, costArr[next.u]));
				}
			}
		}

        PriorityQueue<Item> temp = new PriorityQueue<>();
        while(!itemPq.isEmpty()) {
            Item item = itemPq.poll();
            Item newItem = new Item(item.id, item.revenue, item.dest);
            temp.add(newItem);
        }
        itemPq = temp;
    }

    static private void cancelItem(int id) {
        if(itemMap.get(id) == null) {
            return;
        }
        itemMap.remove(id);
    }

    static private void sell() {
        if(itemPq.isEmpty()) {
            System.out.println(-1);
            return;
        }

        Item item = itemPq.poll();
        while(!itemPq.isEmpty() && itemMap.get(item.id) == null) {
            item = itemPq.poll();
        }
        if(itemMap.get(item.id) == null) {
            System.out.println(-1);
            return;
        }

        if(item.cost < 0) {
            System.out.println(-1);
            itemPq.add(item);
        }else {
            itemMap.remove(item.id);
            System.out.println(item.id);
        }
    }
}

전체 소요 시간 : 3시간

조건을 생각할 게 많아서 정말 까다롭게 풀었던 문제였다.