[백준 20295] 사탕 배달

문제

www.acmicpc.net/problem/20295

 

알고리즘

HLD

 

풀이

$N$개의 가게가 트리 형태로 주어집니다. 가게마다 사탕을 5종류 중 1가지를 팔고 $Q$명의 친구들에게 사탕을 갖다 줄 수 있는지 묻는 문제입니다. 

 

이 문제와 비슷합니다. Milk visits에서는 가게마다 사탕을 2가지 종류 중 1가지를 팔고 있는 것과 마찬가지라고 할 수 있습니다. 이 때는 인접한 노드가 같은 종류의 사탕을 판다면 하나의 집합으로 유니온하는 식으로 풀었습니다. 

 

하지만 현재 문제에서는 5가지 종류가 주어지므로 유니온 파인드 만으로는 구분하기 힘듭니다. 문제를 간단히 하자면 트리 위에 두 정점이 주어졌을 때, 두 정점을 잇는 경로중에서 특정 사탕을 포함하는지를 찾아야 합니다. 비트마스크로 구간마다 어떠한 사탕을 가지고 있는지를 세그먼트 트리위에 저장하면 됩니다.

 

#include <bits/stdc++.h>
#define rep(i, n) for (int i = 0; i < n; ++i)
#define REP(i, n) for (int i = 1; i <= n; ++i)
using namespace std;

const int MAXN = 1e5;
#define left (i << 1)
#define right (i << 1 | 1)

int store[MAXN], p[MAXN], st[MAXN], top[MAXN], idx[MAXN], tree[MAXN << 1],Candy[5];
int N, cnt, Q;
vector<vector<int>> adj;

void dfs(int here) {
    st[here] = 1;
    for (auto& there : adj[here]) {
        if (there ^ p[here]) {
            p[there] = here;
            dfs(there);
            st[here] += st[there];
            int& HE = adj[here][0];
            if (HE == p[here] || st[there] > st[HE])
                swap(HE, there);
        }
    }
}

void hld(int here) {
    idx[here] = cnt++;
    for (auto there : adj[here]) {
        if (there ^ p[here]) {
            top[there] = there == adj[here][0] ? top[here] : there;
            hld(there);
        }
    }
}

int query(int l, int r) {

    int ret = 0;
    l += N;
    r += N;
    for (; l <= r; l >>= 1, r >>= 1) {
        if (l & 1) {
            ret |= tree[l++];
        }
        if (!(r & 1)) {
            ret |= tree[r--];
        }
    }
    return ret;
}

int solve(int u, int v, int candy) {
    int sum = 0;
    while (top[u] ^ top[v]) {
        if (st[top[u]] < st[top[v]])
            swap(u, v);
        sum |= query(idx[top[v]], idx[v]);
        v = p[top[v]];
    }
    if (idx[u] < idx[v])
        swap(u, v);
    
    sum |= query(idx[v], idx[u]);
    return sum & (1 << candy);
}

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

    cin >> N;
    adj.resize(N);
    rep(i, N) {
        cin >> store[i];
        store[i]--;
        Candy[store[i]] = 1;
    }

    rep(i, N - 1) {
        int u, v;
        cin >> u >> v;
        --u, --v;
        adj[u].emplace_back(v);
        adj[v].emplace_back(u);
    }
    dfs(0);
    hld(0);
    rep(i, N) {
        tree[idx[i] + N] = (1<<store[i]);
    }
    for (int i = N - 1; i; --i) {
        tree[i] = tree[left] | tree[right];
    }
    cin >> Q;

    int prv = -1;

    while (Q--) {
        int here, candy;
        cin >> here >> candy;
        --here;
        --candy;
        if (prv == -1 ){
            if(Candy[candy])cout << "PLAY\n";
            else cout << "CRY\n";
        }
        else {
            if (solve(here, prv, candy)) {
            cout << "PLAY\n";
            }
            else {
            cout << "CRY\n";
            }
        }

        prv = here;
    }
    return 0;
}