취미로 코딩하기

※ 글쓴이는 취미로 코딩을 익혀보는 사람이라 정확하지 않은 내용을 담고 있을 수 있다 ※

이번에 볼 문제는 백준 23893번 문제인 알프스의 힘이다.

문제는 아래 링크를 확인하자.

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

편의상 Ai = a, Aj = b라고 표기하자. 문제 조건에 따라 i와 j가 다르다면 a와 b는 항상 서로 다르다.

a와 b가 서로 다르다는 전제 하에서, a^2 + ab + b^2와 k가 합동이라는 것은 (a-b)(a^2 + ab + b^2)와 k(a-b)가 합동이라는 것과 동치관계이다. (a-b의 multiplicative modular inverse가 항상 존재하기 때문이다.)

저 식을 정리하면, 이 문제는 a^3 - ka 와 b^3 - kb 가 서로 합동인 a와 b를 찾는 문제라는 것을 알아낼 수 있다.

오버플로우에 유의하여 위의 관계식을 만족하는 a와 b의 쌍을 잘 구해보자.

입력 순서와 관련있는 i와 j의 순서(i<j 조건)는 크게 신경쓸 필요가 없다. (a,b)와 (b,a)의 두 페어 중 한 가지만을 세기만 하면 같은 값을 구할 수 있기 때문이다.

아래는 제출한 소스코드이다.

#include <iostream>
#include <algorithm>
using namespace std;
typedef long long ll;

ll arr[100000];

int main() {
	ios::sync_with_stdio(0);
	cin.tie(0);
	
	ll N, P, K; cin >> N >> P >> K;
	for (int i = 0; i < N; i++) {
		ll x; cin >> x;
		arr[i] = (((x * x) % P) * x - K * x) % P;
		if (arr[i] < 0) arr[i] += P;
	}
	sort(arr, arr + N);

	ll ans = 0;
	ll old = -1, cnt = 0;
	for (int i = 0; i < N; i++) {
		if (arr[i] == old) cnt++;
		else {
			ans += cnt * (cnt - 1) / 2;
			old = arr[i];
			cnt = 1;
		}
	}
	ans += cnt * (cnt - 1) / 2;

	cout << ans;
}