프로그래머스 | 시소 짝궁

출처

🌐 코딩테스트 연습 - 시소 짝꿍 | 프로그래머스 스쿨

문제

문제 설명

어느 공원 놀이터에는 시소가 하나 설치되어 있습니다. 이 시소는 중심으로부터 2(m), 3(m), 4(m) 거리의 지점에 좌석이 하나씩 있습니다.

이 시소를 두 명이 마주 보고 탄다고 할 때, 시소가 평형인 상태에서 각각에 의해 시소에 걸리는 토크의 크기가 서로 상쇄되어 완전한 균형을 이룰 수 있다면 그 두 사람을 시소 짝꿍이라고 합니다. 즉, 탑승한 사람의 무게와 시소 축과 좌석 간의 거리의 곱이 양쪽 다 같다면 시소 짝꿍이라고 할 수 있습니다.

사람들의 몸무게 목록

weights

이 주어질 때, 시소 짝꿍이 몇 쌍 존재하는지 구하여 return 하도록 solution 함수를 완성해주세요.

제한사항

• 2 ≤ weights의 길이 ≤ 100,000
• 100 ≤ weights[i] ≤ 1,000
  • 몸무게 단위는 N(뉴턴)으로 주어집니다.
  • 몸무게는 모두 정수입니다. 

입출력 예시

weights result

[100,180,360,100,270]

4

풀이

1차 풀이 (실패, 시간초과)

#include <string>
#include <vector>
#include <unordered_map>
#include <set>
using namespace std;

long long solution(vector<int> weights) 
{
    unordered_map<long long, vector<int>> siso_weight_map;
    set<pair<int,int>> pairs;

    for(int i = 0; i < weights.size(); i++)
    {
        for(int j = 0; j < siso_weight_map[weights[i] * 2].size(); j++)
            pairs.insert({i, siso_weight_map[weights[i] * 2][j]});

        for(int j = 0; j < siso_weight_map[weights[i] * 3].size(); j++)
            pairs.insert({i, siso_weight_map[weights[i] * 3][j]});
        
        for(int j = 0; j < siso_weight_map[weights[i] * 4].size(); j++)
            pairs.insert({i, siso_weight_map[weights[i] * 4][j]});
        
        siso_weight_map[weights[i] * 2].push_back(i);
        siso_weight_map[weights[i] * 3].push_back(i);
        siso_weight_map[weights[i] * 4].push_back(i);
    }
        
    return pairs.size();
}

2차 풀이(성공)

#include <string>
#include <vector>
#include <unordered_map>
#include <set>
using namespace std;

long long solution(vector<int> weights) 
{
    unordered_map<long long, vector<int>> siso_weight_map;
    unordered_map<int,long long> weights_map;
    long long counts = 0;    
    
    for(int i = 0; i < weights.size(); i++)
    {
        weights_map[weights[i]] += 1;

        counts += siso_weight_map[weights[i] * 2].size();
        counts += siso_weight_map[weights[i] * 3].size();
        counts += siso_weight_map[weights[i] * 4].size();
        
        siso_weight_map[weights[i] * 2].push_back(i);
        siso_weight_map[weights[i] * 3].push_back(i);
        siso_weight_map[weights[i] * 4].push_back(i);
    }
    
    for(const auto wm : weights_map)
    {
        if(wm.second > 1)
            counts -= wm.second * (wm.second-1); // 중복 
    }
        
    return counts;
}