[BOJ 20532] 정점 간 통신 네트워크

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

20532번: 정점 간 통신 네트워크

Tag : tree dp, number theory

 

문제요약

조상과 자손 관계인 정점 쌍(x, y) 중 a[x]와 a[y]가 약수와 배수 관계인 것의 개수를 구해라.

 

풀이

dfs가 어떻게 돌아가는지 잘 알고 있다면 쉽게 풀 수 있다. 

자손의 입장에서 조상을 세주도록 하겠다. 

일단 루트부터 자신까지의 경로에서 숫자들의 등장 횟수를 구하는 방법은 다음과 같다. 

 

현재 정점의 dfs가 시작되었을 때 cnt[a[x]]을 1 증가시키고 dfs가 끝날 때 cnt[a[x]]를 1 감소시키면

현재 정점의 자손들에서만 cnt[a[x]]가 증가된 상태로 유지된다. 

 

조상 중 자신의 약수를 세는 것은 약수들을 순회하며 cnt[factor]의 합을 구하면 된다. 

배수는 어떻게 구해야 할까? 자손에서 조상이 배수이면 조상에서는 자손이 약수이다. 

조상에서 dfs를 들어가면서 약수들의 개수를 늘리고 끝날 때 감소시키면

자손의 입장에서는 배수들의 개수가 각 b[a[x]]에 담겨있다.

 

이것을 잘 구현하면 된다. 

* a[x]는 a[x]의 약수이기도 하고 배수이기도 하다. 중복을 잘 빼야한다. 

 

+ https://blog.naver.com/jinhan814/222612457695 다행히 진한님이 맞은사람 1등을 탐내지 않으셨다..

 

전체코드

#include "bits/stdc++.h"
# pragma GCC optimize ("O3") 
# pragma GCC optimize ("Ofast") 
# pragma GCC optimize ("unroll-loops") 
# pragma GCC target("sse,sse2,sse3,ssse3,sse4,avx,avx2")
#define endl '\n'

using namespace std;
using ll = long long;

int N, idx;
int a[101010];
int b[101010];
int cnt[101010];
vector<int> factor[101010];
vector<int> v[101010];
ll ans = 0;

void dfs(int x) {
    ans += b[a[x]] - cnt[a[x]];
    for(const auto &i : factor[a[x]]) {
        b[i]++;
        ans += cnt[i];
    }
    cnt[a[x]]++;
    for(const auto &i : v[x]) dfs(i);
    for(const auto &i : factor[a[x]]) b[i]--;
    cnt[a[x]]--;
}

#include <unistd.h>
#include <sys/stat.h>
#include <sys/mman.h>
int z[36];
char*c=(char*)mmap(0,z[12],1,2,0,fstat(0,(struct stat *)z));
inline int f(){int x=0;bool e;c+=e=*c=='-';while(*c>='0')x=10*x+*c++-'0';c++;return e?-x:x;}

int main() {
    ios::sync_with_stdio(0); cin.tie(0);
    
    N = f();
    for(int i=1; i<=N; i++) a[i] = f();
    for(int i=2; i<=N; i++) v[f()].push_back(i);
    for(int i=1; i<=N; i++) for(int j=i; j<=N; j+=i) factor[j].emplace_back(i);
    dfs(1);
    cout << ans << endl;
}