본문 바로가기
알고리즘/코드포스

cf 596 div2 D - Power products (소인수분해, 수학)

by sun__ 2020. 3. 15.

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

 

<문제설명>

최대 1e5의 크기가 n인 배열 a가 주어진다. 임의의 자연수 x에 대해 ai * aj = x^k를 만족하는 쌍의 개수를 구하라.

 

 

<풀이> 0~i-1번 중 i번과 쌍을 이루는~

ai를 소인수분해해서 (j번째로 작은 소수가 나온 횟수 % k)값을 적절히 벡터(v)에 담으면, 0~i-1번의 벡터 중 v와 쌍을 이루는 벡터의 모양은 하나로 정해진다. map<vector<int> int>로 그값을 유지하여 계산해 나간다.

 

1e5이하 소수는 10000개정도 되니까 위 방법을 그대로 적용하려면 시공간복잡도가 너무 크다.

 

 

제곱근1e5 이하의 소수를 smallprimes라 하고 그 외 소수를 bigprime이라고 하자.

 

a는 최대 1e5이므로 제곱근1e5보다 큰 소수(bigprime)는 최대 한 개 갖는다.

-> a는 최대 66개의 smallprimes와 한개의 bigprime을 갖는다.

 

따라서 k>=3인 경우 bigprime을 갖는 ai는 어떤 쌍과도 x^k를 만들지 못한다.

k==2인 경우는 bigprime까지 같은 요소의 개수를 생각해줘야 한다.

-> map<pair<vector<int>, int>, int>로 값을 유지해야 하는 이유.

 

<코드>

int n, k, a[MAX];
map<pair<vector<int>,int>, int> mp;
vector<int> prime; //p^2<=1e5인 소수들
bool isPrime[MAX];

void makeprime() {
	fill(isPrime, isPrime + MAX, true);
	isPrime[1] = false;
	for (ll i = 2; i <= 1e5; i++) {
		if (isPrime[i]) {
			if (i * i <= 1e5) prime.push_back(i);
			for (ll j = i * i; j <= 1e5; j += i)
				isPrime[j] = false;
		}
	}
}

int main() {
	FAST; cin >> n >> k;
	ll xcnt = 0;
	for (int i = 0; i < n; i++) cin >> a[i];
	makeprime();

	int sz = prime.size();
	ll ans = 0;
	for (int i = 0; i < n; i++) {
		vector<int> v(sz), t(sz);
		for (int j = 0; j < sz; j++) {
			int p = prime[j];
			int cnt = 0;
			while (a[i] % p == 0) {
				a[i] /= p;
				cnt++;
			}
			v[j] = cnt % k;
			t[j] = (k - v[j]) % k;
		}
		if (a[i] > 1 && k != 2) continue;
		ans += mp[make_pair(t,a[i])];
		mp[make_pair(v,a[i])]++;			
	}
	cout << ans << '\n';	
}

 

<풀이2> i+1~n번 중 i번과 쌍을 이루는~

미리 v를 만들어두고 풀어봤다. 이경우엔 이분탐색으로 개수를 세주었다. 코드가 더러워진다.

 

<코드2>

ll ans = 0;
int n, k, a[MAX], bigp[MAX], pcnt; // pcnt= 루트1e5이하의 소수개수
vector<int> prime, smallp[MAX];
vector<pair<vector<int>, int>> b;
map<int, int> primeidx;
bool isPrime[MAX];

void init() {
	fill(isPrime, isPrime + MAX, 1);
	int cnt = 0;
	isPrime[1] = false;
	for (ll i = 2; i <= MAX - 1; i++) {
		if (isPrime[i]) {
			prime.push_back(i);
			primeidx[i] = cnt++;
			for (ll j = i * i; j <= MAX - 1; j += i)
				isPrime[j] = false;
		}
	}
	pcnt = lower_bound(prime.begin(), prime.end(), 320) - prime.begin();
	for (int i = 0; i < n; i++)
		smallp[i].resize(pcnt);
}
int main() {
	FAST; cin >> n >> k;
	init();
	
	for (int i = 0; i < n; i++) {
		cin >> a[i];
		int temp = a[i];
		for (int j = 2; j * j <= temp; j++) {
			if (temp % j == 0) {
				int cnt = 0;
				while (temp % j == 0) {
					temp /= j;
					cnt++;
				}
				smallp[i][primeidx[j]] = cnt % k;
			}
		}
		if (temp != 1) {
			if (temp >= prime[pcnt]) bigp[i] = temp;
			else smallp[i][primeidx[temp]]++;
		}
	}

	for (int i = 0; i < n; i++) 
		b.push_back({ smallp[i], bigp[i] });
	sort(b.begin(), b.end());

	for (int i = 0; i < n-1; i++) {
		vector<int> target = b[i].first;
		if (k > 2 && b[i].second) continue;
		for (int j = 0; j < pcnt; j++)
			target[j] = (k - target[j]) % k;
	
		ans += upper_bound(b.begin()+i+1, b.end(), make_pair(target, b[i].second))
			- lower_bound(b.begin()+i+1, b.end(), make_pair(target, b[i].second));
	}
	cout << ans << '\n';
}