배열이 하나 주어지는데 (i, j) : nums[i] == nums[j] 이면서 i<j인 (i, j) 쌍의 개수를 찾아달라고 합니다.

배열 크기가 작으니 아래처럼 그냥 O(n^2)으로 이중포문 쓰면 됩니다.

for (int i=0; i<sz-1; ++i)
    for (int j=i+1; j<sz; ++j)
        if (nums[i] == nums[j])
            ++cnt;

그런데 좀 생각해보면 사실 수가 나오는 순서는 답을 구하는데 연관이 없습니다. 각 수가 몇 번 나오는지에 따라 쌍의 개수가 결정됩니다.

각 수가 n번 나왔다면 이때 문제 조건을 만족하는 쌍의 개수는 \(_nC_2 = n(n-1) / 2\)입니다.

int numIdenticalPairs(vector<int>& nums) {
    unordered_map<int, int> um;
    for (const auto& c : nums) {
        if (um.find(c) != um.end())
            ++um[c];
        else
            um[c] = 1;
    }
    int cnt = accumulate(um.begin(), um.end(), 0, [](int sum, const pair<int, int>& p) {
        const auto& [k, v] = p;
        return sum + (v*(v-1))/2;
    });
    return cnt;
}

풀어보면 소스는 위와 같습니다. std::unordered_map과 std::accumulate를 그냥 써봤습니다.

참고로 accumulate는 left fold를 한다고 합니다.

반응형