[백준 28250번] 이브, 프시케 그리고 푸른 MEX의 아내 (C++)

문제링크 : https://www.acmicpc.net/problem/28250

28250번: 이브, 프시케 그리고 푸른 MEX의 아내

#include <bits/stdc++.h>
using namespace std;
typedef long long ll;
typedef pair<int, int> pii;
const int MAX = 987654321;

int arr[200002];
ll cnt[3];

int main()
{
    ios_base::sync_with_stdio(0); 
	cin.tie(0);
    
    int N;
    cin >> N;
    for(int i=0; i<N; i++)
    {
        cin >> arr[i];
        if(arr[i]==0) cnt[0]++; //0의 갯수
        else if(arr[i]==1) cnt[1]++; //1의 갯수
        else cnt[2]++; //2이상 갯수
    }

    cout << cnt[0] * (cnt[0]-1)/2 + cnt[0]*cnt[1]*2 + cnt[0]*cnt[2]; //0끼리 만남, 0과 1만남, 0과 2이상 만남

    return 0;
}

문제를 살펴보면 아래와 같은 내용을 파악할 수 있다.

Ai = 1, Aj = 1 : MEX(Ai, Aj) = 0
Ai = 1, Aj = 0 : MEX(Ai, Aj) = 2
Ai = 0, Aj = 0 : MEX(Ai, Aj) = 1
Ai = 2, Aj = 0 : MEX(Ai, Aj) = 1

그러면 Ai와 Aj가 모두 0인 경우는 0의 갯수를 m이라 했을 때 mC2인 것을 알 수 있다.

Ai가 0, Aj가 1일때는 Ai의 갯수 * Aj의 갯수에 MEX(Ai, Aj)의 값이 2이기에 추가로 2를 곱해준다.

Ai가 2, Aj가 0일때는 Ai의 갯수 * Aj의 갯수를 곱해주면 된다.

Ai가 1, Aj가 1일때는 마찬가지로 Ai의 갯수 * Aj의 갯수이나 MEX(Ai, Aj)의 값이 0이기에 더해줄 필요가 없다.

 

따라서 위 3가지만 (mC2, cnt[0]*cnt[1]*2 , cnt[0]*cnt[2])의 경우만 더해서 출력해주면 된다.