[BOJ 13059] - Tourists (최소 공통 조상, 트리, 희소 배열, C++, Python)

보양쿠·2023년 7월 4일
0

BOJ

목록 보기
149/260
post-custom-banner

BOJ 13059 - Tourists 링크
(2023.07.04 기준 P5)

문제

n개의 도시와 모든 도시가 연결되어 있게끔 n-1개의 양방향 도로가 주어진다. 도시에는 1부터 n까지의 숫자가 부여되어 있다.
모든 도시의 숫자 x와, x의 배수인 모든 y (y > x)의 쌍에 대해 최단 경로에 포함되어 있는 도시의 수를 합해서 출력

알고리즘

LCA

풀이

"n개의 도시와 모든 도시가 연결되어 있게끔 n-1개의 양방향 도로가 주어진다." -> 트리 구조이다.
트리에선 최소 공통 조상을 이용하여 경로 길이를 구할 수 있다.
그림의 두 경우를 보자.
빨강과 파랑의 깊이. 즉, 각 위치까지 가는 경로에서 초록(LCA)까지 가는 경로가 두 번 겹침을 볼 수 있다. 그러므로 i와 j 사이의 거리 = i의 깊이 + j의 깊이 - lca(i, j)의 깊이 * 2
라고 할 수 있다. 이를 모든 (x, y) 쌍에 대해 구해줘서 더하면 된다.

이 문제에선 간선이 아닌 정점 개수이므로 +1을 해주자.

코드

  • C++
#include <bits/stdc++.h>
using namespace std;

typedef long long ll;

const int MAXN = 200000, MAXH = (int)ceil(log2(MAXN));

int n, h, pa[MAXN][MAXH], lv[MAXN];
vector<int> graph[MAXN];

void dfs(int i, int p){
    for (auto j: graph[i]){
        if (j == p) continue;
        pa[j][0] = i;
        lv[j] = lv[i] + 1;
        dfs(j, i);
    }
}

int lca(int i, int j){
    if (lv[i] < lv[j]) swap(i, j);
    int dif = lv[i] - lv[j];

    int k = 0;
    while (dif){
        if (dif & 1) i = pa[i][k];
        dif >>= 1; k++;
    }

    if (i != j){
        for (k = h - 1; k >= 0; k--)
            if (pa[i][k] != pa[j][k]) i = pa[i][k], j = pa[j][k];
        i = pa[i][0];
    }

    return i;
}

int main(){
    ios_base::sync_with_stdio(0);
    cin.tie(0);

    cin >> n;
    for (int _ = 1, i, j; _ < n; _++){
        cin >> i >> j;
        graph[--i].push_back(--j);
        graph[j].push_back(i);
    }

    // 희소 배열
    h = (int)ceil(log2(n));
    fill(&pa[0][0], &pa[n - 1][h], 0);
    fill(lv, lv + n, 0);
    dfs(0, -1);

    for (int j = 1; j < h; j++) for (int i = 0; i < n; i++)
        pa[i][j] = pa[pa[i][j - 1]][j - 1];

    // i와 j 사이의 거리 = i의 깊이 + j의 깊이 - lca(i, j)의 깊이 * 2
    ll result = 0;
    for (int i = 1, half = n / 2; i <= half; i++)
        for (int j = i * 2; j <= n; j += i)
            // 간선의 개수가 아닌 정점의 개수이므로 1을 더하자.
            result += lv[i - 1] + lv[j - 1] - lv[lca(i - 1, j - 1)] * 2 + 1;
    cout << result;
}
  • Python
import sys; input = sys.stdin.readline
sys.setrecursionlimit(3000000)
from math import ceil, log2

def dfs(i, p):
    for j in graph[i]:
        if j == p:
            continue
        pa[j][0] = i
        lv[j] = lv[i] + 1
        dfs(j, i)

def lca(i, j):
    if lv[i] < lv[j]:
        i, j = j, i
    dif = lv[i] - lv[j]

    k = 0
    while dif:
        if dif & 1:
            i = pa[i][k]
        dif >>= 1
        k += 1

    if i != j:
        for k in range(h - 1, -1, -1):
            if pa[i][k] != pa[j][k]:
                i = pa[i][k]
                j = pa[j][k]
        i = pa[i][0]

    return i

n = int(input())
graph = [[] for _ in range(n)]
for _ in range(n - 1):
    i, j = map(int, input().split())
    i -= 1; j -= 1
    graph[i].append(j)
    graph[j].append(i)

# 희소 배열
h = ceil(log2(n))
pa = [[0] * h for _ in range(n)]
lv = [0] * n
dfs(0, -1)

for j in range(1, h):
    for i in range(n):
        pa[i][j] = pa[pa[i][j - 1]][j - 1]

# i와 j 사이의 거리 = i의 깊이 + j의 깊이 - lca(i, j)의 깊이 * 2
result = 0
for i in range(1, n // 2 + 1):
    for j in range(i * 2, n + 1, i):
        # 간선의 개수가 아닌 정점의 개수이므로 1을 더하자.
        result += lv[i - 1] + lv[j - 1] - lv[lca(i - 1, j - 1)] * 2 + 1
print(result)
profile
GNU 16 statistics & computer science
post-custom-banner

0개의 댓글