[SWEA] 4534 트리 흑백 색칠

출처: SWEA

Approach

메모제이션을 이용해서 풀 수 있는 문제이다.

DP[Color][Node] = Node에서 Color를 가질 때, 가질 수 있는 최대 색칠 방법 수

    - DP[흰색][node] = DP[검은색][next] + DP[흰색][next]

    - DP[검은색][node] = DP[흰색][next]

 

변수 계산에서 int 범위를 벗어나기 때문에 long long 타입이 필요하고

 MOD로 나머지 연산(%) 처리해야 한다.

 

▶ 동적계획법(Dynamic Programming, DP) 

동적계획법(Dynamic Programming, DP)

---

#include <stdio.h>
#define LL long long

enum {
    WHITE = 0,
    BLACK = 1,
};
const int MAX_N = 1e5 + 1;
const int MOD = 1000000007;

struct Node
{
    int x, y;
}edge[MAX_N * 2], temp[MAX_N * 2];
int TC, N, ans;
int bound, flag[MAX_N], DP[2][MAX_N];

void input()
{
    scanf("%d", &N);
    for (int i = 0; i <= N; i++)
        DP[WHITE][i] = DP[BLACK][i] = -1;
    for (int i = 0; i < N - 1; i++) {
        scanf("%d %d", &edge[i].x, &edge[i].y);
        edge[i + N - 1].y = edge[i].x;
        edge[i + N - 1].x = edge[i].y;
    }
    bound = 2 * N - 1;
    edge[bound - 1].x = 0;
    edge[bound - 1].y = 1;
}

void mergeSort(int start, int end) {
    // 0. base condition
    if (start >= end) return;
    // 1. divide
    int mid = (start + end) / 2;
    // 2. conquer
    mergeSort(start, mid);
    mergeSort(mid + 1, end);
    // 3. merge
    int i = start, j = mid + 1, k = start;
    while (i <= mid && j <= end) {
        if (edge[i].x < edge[j].x)
            temp[k++] = edge[i++];
        else
            temp[k++] = edge[j++];
    }
    while (i <= mid) temp[k++] = edge[i++];
    while (j <= end) temp[k++] = edge[j++];
    // 4. copy
    for (i = start; i <= end; ++i) {
        edge[i] = temp[i];
    }
}

int upperBound(const int& pivot) {
    int start = 0, end = N * 2 - 2;
    while (start <= end) {
        int mid = (start + end) / 2;
        if (edge[mid].x < pivot)
            start = mid + 1;
        else
            end = mid - 1;
    }
    return start;
}

int DFS(int from, int color) {
    if (~DP[color][from])
        return DP[color][from];
    
    LL ret = flag[from] = 1;
    int cur = upperBound(from);
    
    while (edge[cur].x == from && cur < bound)
    {
        int to = edge[cur].y;
        if (!flag[to]) {
            if (color == WHITE)
            {
                ret *= DFS(to, BLACK) + DFS(to, WHITE);
            }    
            else // BLACK
            {
                ret *= DFS(to, WHITE);
            }
                
            ret %= MOD;
        }
        cur++;
    }
    
    flag[from] = 0;
    DP[color][from] = ret;
    return ret;
}

int main() {
    // freopen("input.txt", "r", stdin);
    scanf("%d", &TC);
    for (int tc = 1; tc <= TC; tc++) {
        input();
        mergeSort(0, bound - 1);
        printf("#%d %d\n", tc, DFS(0, WHITE));
    }
    return 0;
}