백준 20303번

오늘 풀어본 문제는 백준의 20303번 문제¹이다. 문제 풀이에 사용한 언어는 C++ 이다.

solved.ac 기준 CLASS

CLASS 5

문제 정보

이 문제의 내용과 조건은 다음과 같다.

문제

Trick or Treat!!

10월 31일 할로윈의 밤에는 거리의 여기저기서 아이들이 친구들과 모여 사탕을 받기 위해 돌아다닌다. 올해 할로윈에도 어김없이 많은 아이가 할로윈을 즐겼지만 단 한 사람, 일찍부터 잠에 빠진 스브러스는 할로윈 밤을 즐길 수가 없었다. 뒤늦게 일어나 사탕을 얻기 위해 혼자 돌아다녀 보지만 이미 사탕은 바닥나 하나도 얻을 수 없었다.

단단히 화가 난 스브러스는 거리를 돌아다니며 다른 아이들의 사탕을 빼앗기로 마음을 먹는다. 다른 아이들보다 몸집이 큰 스브러스에게 사탕을 빼앗는 건 어렵지 않다. 또한, 스브러스는 매우 공평한 사람이기 때문에 한 아이의 사탕을 뺏으면 그 아이 친구들의 사탕도 모조리 뺏어버린다. (친구의 친구는 친구다?!)

사탕을 빼앗긴 아이들은 거리에 주저앉아 울고 $K$ 명 이상의 아이들이 울기 시작하면 울음소리가 공명하여 온 집의 어른들이 거리로 나온다. 스브러스가 어른들에게 들키지 않고 최대로 뺏을 수 있는 사탕의 양을 구하여라.

스브러스는 혼자 모든 집을 돌아다녔기 때문에 다른 아이들이 받은 사탕의 양을 모두 알고 있다. 또한, 모든 아이는 스브러스를 피해 갈 수 없다.

입력

첫째 줄에 정수 $N$, $M$, $K$ 가 주어진다. $N$ 은 거리에 있는 아이들의 수, $M$ 은 아이들의 친구 관계 수, $K$ 는 울음소리가 공명하기 위한 최소 아이의 수이다. $(1 \leq N \leq 30\ 000$, $0 \leq M \leq 100\ 000$, $1 \leq K \leq \min\left{N, 3\ 000\right})$

둘째 줄에는 아이들이 받은 사탕의 수를 나타내는 정수 $c_1, c_2, \cdots, c_N$이 주어진다. $(1 \leq c_i \leq 10\ 000)$

셋째 줄부터 $M$개 줄에 갈쳐 각각의 줄에 정수 $a$, $b$ 가 주어진다. 이는 $a$ 와 $b$ 가 친구임을 의미한다. 같은 친구 관계가 두 번 주어지는 경우는 없다. $(1 \leq a, b \leq N$, $a \neq b)$

출력

스브러스가 어른들에게 들키지 않고 아이들로부터 뺏을 수 있는 최대 사탕의 수를 출력한다.

풀이과정

1번째 시도

문제를 해결하기 위해 생각한 방식은 다음과 같다.

Union-Find 를 이용하여 아이들을 그룹별로 묶는다.
각 그룹의 아이들의 수와 사탕의 총 개수를 각각 children 과 candies 의 같은 인덱스에 저장한다.
children 을 무게로, candies 를 값어치로 보고 배낭 문제를 푼다.

Union-Find 는 기존에 만들어 둔 DisjointSet 을 그대로 이용하였다. 예전 글에서도 계속 사용했으니 이번에는 코드를 생략하겠다.

배낭 문제가 종종 나오는 것 같아 나중에 재활용할 수 있도록 knapsack.hpp 파일을 만들고, knapsack 함수를 거기에 구현해 두었다.

코드는 다음과 같이 작성하였다. 제출할 때는 헤더파일의 코드를 붙여넣어 제출하였다.

knapsack.hpp

#ifndef __KNAPSACK_HPP__
#define __KNAPSACK_HPP__

#include <vector>

template <typename T>
T knapsack(const std::vector<T>& weight, const std::vector<T>& value, const T& maxWeight, size_t currItem = 0) {
    if (weight.size() != value.size()) {
        return static_cast<T>(0);
    }
    else if (currItem >= weight.size() || maxWeight <= 0) {
        return static_cast<T>(0);
    }
    else if (weight[currItem] > maxWeight) {
        return knapsack(weight, value,  maxWeight, currItem + 1);
    }
    else {
        T withoutCurr = knapsack(weight, value, maxWeight, currItem + 1);
        T withCurr = value[currItem] + knapsack(weight, value, maxWeight - weight[currItem], currItem + 1);

        return withoutCurr > withCurr ? withoutCurr : withCurr;
    }
}

#endif //__KNAPSACK_HPP__

main.hpp

#include <bits/stdc++.h>
#include "disjoint_set.hpp"
#include "knapsack.hpp"

using namespace std;

using ll = long long;
using ull = unsigned long long;

using pii = pair<int, int>;
using pll = pair<ll, ll>;

int main(void) {
    ios::sync_with_stdio(false);
    cin.tie(nullptr);
    cout.tie(nullptr);

    int N, M, K;
    cin >> N >> M >> K;

    vector<int> c(N+1, 0);

    for (int i=1; i<=N; i++) {
        cin >> c[i];
    }

    DisjointSet diset(N+1, 0);

    for (int i=0; i<M; i++) {
        int a, b;
        cin >> a >> b;

        diset.merge(a, b);
    }

    vector<int> groupedChildren(N+1, 0);
    vector<int> groupedCandies(N+1, 0);

    vector<int> groupNumbers;

    for (int i=1; i<=N; i++) {
        int curr = diset.find(i);

        if (groupedCandies[curr] == 0) {
            groupNumbers.emplace_back(curr);
        }

        groupedChildren[curr] += 1;
        groupedCandies[curr] += c[i];
    }

    vector<int> candies;
    vector<int> children;

    for (auto& index : groupNumbers) {
        children.emplace_back(groupedChildren[index]);
        candies.emplace_back(groupedCandies[index]);
    }

    cout << knapsack<int>(children, candies, K - 1);

    return 0;
}

실행 결과 ‘시간 초과’ 가 떴다.

2번째 시도

시간 초과가 난 이유로 생각해본 것은 재귀 함수를 호출하는 과정에서 불필요한 호출이 너무 많이 일어나는 것이었다. 특정 조합의 인자를 기반으로 함수가 이미 호출된적이 있는데, 딱히 저장하거나 다시 활용하지 않는 것이었다.

결국 knapsack 함수를 전체적으로 수정하기로 했고, 재귀함수의 형태로 구성하는 대신 2차원 배열에 값을 저장한 뒤, 이전 값을 불러와 현재 값을 계산하도록 하였다.

코드는 다음과 같이 수정하였다. (knapsack.hpp 파일만 수정했다.)

knapsack.hpp

#ifndef __KNAPSACK_HPP__
#define __KNAPSACK_HPP__

#include <vector>

template <typename T>
T knapsack(const std::vector<T>& weight, const std::vector<T>& value, const T& maxWeight) {
    if (weight.size() != value.size()) {
        return static_cast<T>(0);
    }

    size_t n = weight.size();
    std::vector<std::vector<T>> dp(n + 1, std::vector<T>(maxWeight + 1, 0));

    for (size_t i=1; i<=n; i++) {
        for (T w=0; w<=maxWeight; w++) {
            T without = dp[i-1][w];

            if (weight[i - 1] <= w) {
                T with = value[i-1] + dp[i-1][w - weight[i-1]];
                dp[i][w] = with > without ? with : without;
            }
            else {
                dp[i][w] = without;
            }
        }
    }

    return dp[n][maxWeight];
}

#endif //__KNAPSACK_HPP__

그러자 모든 테스트 케이스를 통과하고 ‘맞았습니다’가 나오는 것을 확인할 수 있었다.

마무리

문제가 꽤 특이했던 것 같다. Union-Find 와 배낭 문제 해결을 동시에 구현하는 경우 처음 봤던 것이다. 다행히 DisjointSet 을 잘 만들어 두어 Union-Find 부분은 문제가 없었지만, 이번에 배낭 문제는 시간 초과의 원인을 생각하고 함수를 구현하느라 힘을 꽤 썼던 것 같다. 이번에는 이 코드를 파일로 저장해두었으니 다음 번에는 더 수월하게 배낭 문제를 풀 수 있을 것이다.

내가 구현한 DisjointSet 의 코드²와 knapsack 함수의 코드³는 내 레포지토리에서 확인할 수 있다.

오늘의 PS는 여기까지!

1: https://www.acmicpc.net/problem/20303

2: https://github.com/ChoiCube84/baekjoon-solutions/blob/main/C%2B%2B/custom_data_structures/disjoint_set.hpp

3: https://github.com/ChoiCube84/baekjoon-solutions/blob/main/C%2B%2B/custom_data_structures/knapsack.hpp

← Previous Post Next Post →