Algorithm

[백준] 1202번 보석 도둑 (javascript, python)

주인장 꼬비 2024. 1. 30. 02:17

그리디 문제를 풀고있는데 또다시 어려움을 직면했고 오답노트겸 따로 정리를 했다. 

 

 

https://www.acmicpc.net/problem/1202

 

1202번: 보석 도둑

첫째 줄에 N과 K가 주어진다. (1 ≤ N, K ≤ 300,000) 다음 N개 줄에는 각 보석의 정보 Mi와 Vi가 주어진다. (0 ≤ Mi, Vi ≤ 1,000,000) 다음 K개 줄에는 가방에 담을 수 있는 최대 무게 Ci가 주어진다. (1 ≤ Ci

www.acmicpc.net

 

문제

세계적인 도둑 상덕이는 보석점을 털기로 결심했다.

상덕이가 털 보석점에는 보석이 총 N개 있다. 각 보석은 무게 Mi와 가격 Vi를 가지고 있다. 상덕이는 가방을 K개 가지고 있고, 각 가방에 담을 수 있는 최대 무게는 Ci이다. 가방에는 최대 한 개의 보석만 넣을 수 있다.

상덕이가 훔칠 수 있는 보석의 최대 가격을 구하는 프로그램을 작성하시오.

 

입력

첫째 줄에 N과 K가 주어진다. (1 ≤ N, K ≤ 300,000)

다음 N개 줄에는 각 보석의 정보 Mi와 Vi가 주어진다. (0 ≤ Mi, Vi ≤ 1,000,000)

다음 K개 줄에는 가방에 담을 수 있는 최대 무게 Ci가 주어진다. (1 ≤ Ci ≤ 100,000,000)

모든 숫자는 양의 정수이다.

 

 

출력

첫째 줄에 상덕이가 훔칠 수 있는 보석 가격의 합의 최댓값을 출력한다.

 

예제 입력 1

2 1
5 10
100 100
11

예제 출력 1

10

예제 입력 2

3 2
1 65
5 23
2 99
10
2

예제 출력 2

164

 


 

최초 시도 (시간 초과로 실패)

const input = require("fs")
  .readFileSync(process.platform === "linux" ? "/dev/stdin" : "../input.txt")
  .toString()
  .trim()
  .split("\n");

function solution(input) {
  let answer = 0;

  const [N, K] = input[0].split(" ").map(Number);
  const MV = input.slice(1, N + 1); // M: 보석 무게, V: 보석 가격
  const C = input.slice(N + 1, N + K + 1).map(Number);

  const sortedC = C.sort((a, b) => b - a); // 무거운거 담을 수 있는 가방부터
  const sortedMV = MV.sort((a, b) => {
    // 비싼 보석 순으로 정렬
    const numA = parseInt(a.split(" ")[1], 10);
    const numB = parseInt(b.split(" ")[1], 10);

    return numB - numA;
  });

  sortedC.map((bagWeight) => {
    for (let idx in sortedMV) {
      let currentMV = sortedMV[idx];
      if (!currentMV) continue;

      const [M, V] = currentMV.split(" ").map(Number);
      if (bagWeight >= M) {
        answer += V;
        sortedMV[idx] = null;
        break;
      }
    }
  });
  return answer;
}

console.log(solution(input));

 

가방은 가벼운 보석을 넣을 수 있는 가방이 앞으로 오게 정렬하고, 보석은 가격이 비싼게 앞으로 오게 정렬을 하는 것으로 시작했다.

그리고 가방 배열을 돌면서 넣을 수 있는 무게 중에 가장 비싼거부터 넣고, 다음 가방에 넣기 전에 해당 값을 Null처리함으로서 같은 보석을 여러번 넣지않게 구현했다. 

시간 초과로 틀리고 나서 다시 문제를 살펴보니 지금 코드로 대략 시간 복잡도를 계산해보니 보석(300,000) * 가방(300,000) = 90,000,000,000 이 나왔다. 

 

이리저리 다른 방법으로 풀어보다가 python 의 heapq (혹은 PriorityQueue)를 사용하면 풀릴 것 같았는데 js로 heapq를 구현하기가 귀찮고 확신이 없어서 결국 다른 사람의 소스코드를 찾아봤다.

대충 2~3시간 삽질했달까

 

 

다음 코드는 주석에 달린 블로그에서 가져와서 좀 만져보고 연구해보고 살짝 리팩토링한 코드다.

/**
 * https://lhoiktiv.tistory.com/580
 */
class MaxHeap {
  constructor() {
    this.heap = [];
  }
  swap(a, b) {
    // 구조분해 할당 문법으로 swap 가능
    [this.heap[a], this.heap[b]] = [this.heap[b], this.heap[a]];
  }
  size() {
    return this.heap.length;
  }
  push(value) {
    // 맨뒤에 추가 max heap 이므로 부모랑 비교해서 큰값을 부모랑 swap 해줘야함
    this.heap.push(value);
    let idx = this.heap.length - 1;
    let parent = Math.floor((idx - 1) / 2);

    while (this.heap[parent] < value) {
      this.swap(parent, idx);
      idx = parent;
      parent = Math.floor((idx - 1) / 2);
    }
    // return this.print()
  }
  // 큐이기 때문에 삭제는 항상 루트노드부터 이루어짐. 루트 노드를 삭제하고, 맨마지막 인덱스를 루트랑 교환

  pop() {
    const lastIdx = this.heap.length - 1;
    let idx = 0;
    this.swap(0, lastIdx); // 0번이 루트노드
    let value = this.heap.pop();

    while (idx < lastIdx) {
      let leftChildIdx = idx * 2 + 1;
      let rightChildIdx = idx * 2 + 2;

      // 왼쪽자식 인덱스가 더 크다는 뜻은 자식노드가 없다는 뜻
      if (leftChildIdx >= lastIdx) {
        break;
      } else if (rightChildIdx >= lastIdx) {
        // 왼쪽 자식만 있는경우 자식과 비교해서 크면 스왑
        if (this.heap[idx] < this.heap[leftChildIdx]) {
          this.swap(idx, leftChildIdx);
          idx = leftChildIdx;
        } else {
          break;
        }
      } else {
        // 둘다 있는경우 중 두 자식이 루트보다 다 큰경우
        if (this.heap[leftChildIdx] > this.heap[idx] && this.heap[rightChildIdx] > this.heap[idx]) {
          // 큰값이랑 스왑
          if (this.heap[leftChildIdx] > this.heap[rightChildIdx]) {
            this.swap(idx, leftChildIdx);
            idx = leftChildIdx;
          } else {
            this.swap(idx, rightChildIdx);
            idx = rightChildIdx;
          }
        } else if (this.heap[leftChildIdx] > this.heap[idx]) {
          // 왼쪽 자식만 루트보다 클 경우
          this.swap(leftChildIdx, idx);
          idx = leftChildIdx;
        } else if (this.heap[rightChildIdx] > this.heap[idx]) {
          // 오른쪽 자식
          this.swap(rightChildIdx, idx);
          idx = rightChildIdx;
        } else {
          // 둘다 작을경우 안바꿈
          break;
        }
      }
    }
    return value;
  }
}

const input = require("fs")
  .readFileSync(process.platform === "linux" ? "/dev/stdin" : "../input.txt")
  .toString()
  .trim()
  .split("\n")
  .map((v) => v.split(" ").map(Number));

function solution(input) {
  let answer = 0;
  const heap = new MaxHeap();

  const [N, K] = input[0]; // N: 보석의 수, K: 가방의 수

  // 무게가 가벼운 보석이 앞으로 오게 정렬
  let jewel = input.slice(1, N + 1).sort((a, b) => a[0] - b[0]);

  // 가벼운 걸 담을 수 있는 가방이 앞으로 오게 정렬
  let bag = input
    .slice(N + 1, N + 1 + K)
    .map((v) => v[0])
    .sort((a, b) => a - b);

  let j = 0;
  for (let i = 0; i < K; i++) {
    while (j < N && jewel[j][0] <= bag[i]) {
      heap.push(jewel[j][1]);
      j++;
    }
    if (heap.size()) {
      answer += heap.pop();
    }
  }
  return answer;
}

console.log(solution(input));

 

간단히 설명하자면

보석은 우선 무게가 가벼운게 앞으로 오게 정렬하였고,

가방은 가벼운 것을 담을 수 있는게 앞으로 오게 정렬하였다.

 

그리고 for문으로 가벼운 가방부터 시작하여 현재 가방에 담을 수 있는 무게 이하의 보석들을 MaxHeap에 추가한다.

현재 가방에 담길 수 있는 최대 무게의 보석까지 담기면 heap.pop()으로 Maxheap에 있는 보석중에 가장 가치가 있는 보석을 가져와서 answer에 더해준다.

 

이 과정을 모든 보석(j)을 사용할때 까지 && 가방이 수용할 수 있는 최대 무게의 보석까지 처리한 경우

까지 반복한다음 answer를 리턴하는 방식이다.

 

시간 복잡도는 HeapPush( O(NlogN) ) + HeapPop( O(KlogN) ) 라서 전체 시간 복잡도는 (N+K)logN 이라고 볼 수 있다.

 

 

 


 

 

<2024.02.03 내용 추가>

파이썬으로 알고리즘 문제푸는게 정신건강에 좋겠다는 생각이 들어 파이썬으로 이 문제를 다시 풀어보았다.

입력은 한꺼번에 받는 것이 아니라 줄 단위로 입력받게 하였고,

heapq를 이용하여 풀었다. (PriorityQueue 로 풀 수도 있지만 코테 환경에서는 보통 heapq가 더 빠른걸로 알고 있어서 heapq로 풀었다.) heapq는 MinHeap만 제공되기 때문에 해당 문제를 풀기위해 `-` 를 앞에 붙여줌으로서 MaxHeap과 동일한 결과를 얻을 수 있게 하였다. 

 

import heapq
import sys

input=sys.stdin.readline

def solution():
  n, k = map(int, input().split())

  gems = []
  for i in range(1, n+1):
      gems.append(list(map(int, input().split())))
  gems.sort() # 가벼운 보석 순이 앞으로 오게 정렬

  bags = []
  for i in range(n+1, n+1+k):
      bags.append(int(input()))
  bags.sort() # 담을 수 있는 보석의 무게가 가벼운것부터 가방 정렬

  result = 0
  tmp = []
  i = 0
  for bag in bags:
      while i < len(gems) and gems[i][0] <= bag:
          heapq.heappush(tmp, -gems[i][1])
          i+=1
      if len(tmp):
          result -= heapq.heappop(tmp)
  return result

print(solution())