[Programmers] Level 2 :: 시소 짝꿍 C++ 풀이

프로그래머스

N개의 몸무게가 주어지고, 시소가 주어진다. 시소에는 2, 3, 4 칸 중 하나에 앉을 수 있다.

이때 균형을 맞출 수 있도록 앉는 경우의 수를 구해야 한다. N은 10만 이하이므로 O(N^2)에는 안풀린다.
우리가 궁금한 것은 개수다. 그러니 각 몸무게별로 인원 수를 map에 저장하자. 그리고 이 key 값들을 이용해서 짝의 개수를 세면 된다.

앉을 수 있는 자리가 2, 3, 4이므로 2중 for문을 돌려주자. 현재 몸무게를 w라고 할때, 2, 3, 4 모두에 앉혀봐야 한다. 그리고 이때 반대편에 앉을 수 있는 몸무게가 존재한다면 개수를 곱해주면 된다. 

반대편에 앉을 수 있는 몸무게는 w * i / j와 같은 수식으로 구할 수 있다. 만약 (w * i) % j == 0이 아니라면 애초에 반대편 몸무게가 존재하지 않으므로, 이를 만족하는 w를 기준으로만 세면 된다.

#include <string>
#include <vector>
#include <map>
#include <iostream>
#include <set>
#include <cstring>
using namespace std;
set<int> st;
map<int, long long> mp;
bool visited[1010][1010];

long long solution(vector<int> weights) {
    memset(visited, false, sizeof(visited));
    long long answer = 0;
    for (int weight : weights) {
        mp[weight]++;
        st.insert(weight);
    }
    for (auto w : st) {
        answer += (mp[w] * (mp[w]-1)) / 2;
        for (int i=2; i<=4; i++) {
            for (int j=2; j<=4; j++) {
                if (i == j) continue;
                if ((w * i) % j) continue;
                if (mp[w*i/j] != 0 && !visited[w*i/j][w]) {
                    visited[w*i/j][w] = true;
                    visited[w][w*i/j] = true;
                    answer += mp[w] * mp[w*i/j];
                }
            }
        }
    }
    return answer;
}