Educational #56 D - GCD counting (트리(숲) 지름, 소인수분해)

https://codeforces.com/contest/1101/problem/D

소인수분해 시간복잡도를 잘못 파악하고 있었다. 루트n인줄 알았는데 소인수분해는 loglogn이다.

(에라토스가 루트n이다)

 

트리 지름 구하는 방법은 알고 있었지만 숲에서 최대 지름을 구하는 것을 O(n)에 할 방법을 생각하지 못했다.

 

같은 시간복잡도를 갖는 코드라도 시간제한이 빡빡해서 최적화를 해줘야 한다. (난이도에 비해 푼 사람이 적은 이유인듯)

 

 

<숲에서 최대 지름 O(nlogn)에 구하기>

단일 트리와 다른 특별한 알고리즘을 사용하진 않는다. 한 점 잡고 제일 먼 점 u와 u에서 제일 먼 점 v와의 거리를 모든 트리에 대해 구하면 된다.

 

하나의 트리에 대해 최대 지름을 구한 후, 방문처리를 위한 vst배열을 초기화하지 않고 다음 트리에 대해 최대 지름을 구하는 부분에 주목하자.

 

<코드>

bool used[MAX]; //트리 구분용
int counter, diameter, vst[MAX], d[MAX];

P bfs(int st) { // {st에서 가장 먼 점, 거기까지 거리} 반환
	counter++;
	queue<int> q; q.push(st);
	vst[st] = counter;
	d[st] = 1;
	P ret = { st,d[st] };
	while (!q.empty()) {
		int u = q.front(); q.pop();
		used[u]=true;
		ret = { u,d[u] };

		for (int v : adj[u]) {
			if (vst[v] != counter) {
				vst[v] = counter;
				d[v] = d[u] + 1;
				q.push(v);
			}
		}
	}
	return ret;

}

int main() {

	....

	memset(used,0, sizeof(used))
	for (int u = 1; u <= n; u++) if (!used[u]) {
		P vd = bfs(u);

		//현재 트리에 속한 정점만 방문배열 초기화 해야함.
		//->bfs를 함수로 구현하고 counter로 초기화를 대신함.
		//위 방법을 사용하지 않으면 이 부분에 memset(vst,0,sizeof(vst))들어가서 n제곱이 될 것.
		P vd2 = bfs(vd.first);
		diameter = max(diameter, vd2.second);
	}

	cout << diameter << '\n'; //숲에서 지름.
	....
}

 

---

 

<문제설명>

최대 1e5개의 정점 n개로 이뤄진 트리가 주어진다. 각 정점마다 값(a)을 따로 갖는다.

 

트리의 최대 길이를 구하되, 경로상 모든 정점의 gcd값이 1 이상인 최대 길이를 구하라.

 

 

<풀이>

200000이하의 모든 소수 p에 대해 p로 나눠지는 정점들에 대해서만 숲의 지름을 구해주며 최댓값을 갱신해가면 된다.

 

-> 전처리 후 O((200000이하소수개수) * n) -> 대충O(2e9) 

 

소수p로 나눠지는 정점을 미리 벡터 dbp에 넣어두고 그 정점만 사용해서 숲의 지름을 구해주면 시간안에 겨우 돌아간다.

 

 

<코드>

int n, a[MAX], ans, counter, vst[MAX], d[MAX];
vector<int> adj[MAX], dbp[MAX], g[MAX];
set<int> primes, used;

P mxvertDistFrom(int st) {
	counter++;
	queue<int> q;
	q.push(st);
	d[st] = 1;
	vst[st] = counter;
	P ret = { st,d[st] };
	while (!q.empty()) {
		int u = q.front(); q.pop();
		used.insert(u);

		ret = { u, d[u] };
		for (int v : g[u]) {
			if (vst[v] != counter) {
				vst[v] = counter;
				d[v] = d[u] + 1;
				q.push(v);
			}
		}
	}
	return ret;
}

void go(int p) {
	if (dbp[p].empty()) return;

	for (int u : dbp[p]) {
		g[u].clear();
		for (int v : adj[u]) if (a[v] % p == 0)
			g[u].push_back(v);
	}

	memset(used, 0, sizeof(used));
	for (int u : dbp[p]) if (!used[u]) {
		P vd = mxvertDistFrom(u);
		P vd2 = mxvertDistFrom(vd.first);
		ans = max(ans, vd2.second);
	}
}

int main() {
	FAST; cin >> n;
	for (int i = 1; i <= n; i++) {
		cin >> a[i];
		int t = a[i];
		for (ll j = 2; j * j <= t; j++) {
			if (t % j == 0) {
				primes.insert(j);
				dbp[j].push_back(i);
				while (t % j == 0) t /= j;
			}
		}
		if (t > 1) {
			primes.insert(t);
			dbp[t].push_back(i);
		}
	}
	for (int i = 0, u, v; i < n - 1; i++) {
		cin >> u >> v;		
		adj[u].push_back(v);
		adj[v].push_back(u);	
	}

	for (int i : primes) go(i);
	cout << ans << '\n';
}

 

<생각>

내가 잘못 생각한 것이 아니라면, go함수 내 memset(used,0,sizeof(used))이후 반복문에

 

for(int u=1; u<=n;u++)로 바꿔도 시간복잡도는 같다. 하지만 이렇게 바꾸면 TLE가 나온다.

 

휴리스틱문제도 아니고 같은 시간복잡도를 갖는 코드가 어떤것은 맞고 어떤것은 틀리다는 것은 문제가 잘못된것 아닌가 생각한다.

 

 

숲에서 최대 지름구하는 부분은 유념하자.