C++で配列内のK-diffペアを求める方法

問題の概要

配列と整数kが与えられたとき、配列内に存在する「一意なk-diffペア」の数を求めることを考えます。ここでk-diffペアとは (i, j) のように表されるペアであり、i と j がどちらも配列内に存在し、両者の絶対差が k と一致するものを指します。

例えば、入力が [3,1,4,1,5]、k = 2 の場合を考えてみましょう。このとき配列には (1,3) と (3,5) という2つの2-diffペアが存在するため、出力は 2 となります。

解決のためのアプローチ

この問題は、マップとセットを組み合わせることで効率的に解くことができます。手順は以下の通りです。

• seen と done という2つのマップを定義します。
• セット s を1つ定義します。
• k < 0 の場合は 0 を返します（絶対差が負になることはないため）。
• i := 0 で初期化し、nums のサイズ未満である間 i を1ずつ増やしながら以下を繰り返します。
  • seen[nums[i]] の値を1増やします。
  • nums[i] をセット s に挿入します。
• ans := 0 で初期化します。
• セット s 内の各要素 it に対して以下を処理します。
  • k が 0 と等しい場合：
    • seen[it] > 1 であれば（同じ値が2回以上出現していれば）、ans を1増やします。
  • それ以外の場合：
    • done[it] を1増やします。
    • (it + k) が seen に存在し、かつ done に存在しない場合は ans を1増やします。
    • (it - k) が seen に存在し、かつ done に存在しない場合は ans を1増やします。
• 最後に ans を返します。

この手法のポイントは、セット s が要素をソート済みの状態で保持するため、小さい値から順に処理が進み、done マップによって同じペアの重複カウントを確実に防げる点です。また、k = 0 のケースでは「同じ値が2回以上出現しているか」を確認することで、同値同士のペアも正しく数えることができます。

実装例

理解を深めるために、以下のC++による実装例を見てみましょう。

#include <bits/stdc++.h>
using namespace std;
class Solution {
public:
   int findPairs(vector<int>& nums, int k) {
      map<int, int> seen, done;
      set<int> s;
      if (k < 0)
         return 0;
      for (int i = 0; i < nums.size(); i++) {
         seen[nums[i]]++;
         s.insert(nums[i]);
      }
      int ans = 0;
      for (auto it = s.begin(); it != s.end(); it++) {
         if (k == 0) {
            if (seen[*it] > 1)
            ans++;
         }
         else {
            done[*it]++;
            if (seen.find(*it + k) != seen.end() && done.find(*it + k) == done.end())
               ans++;
            if (seen.find(*it - k) != seen.end() && done.find(*it - k) == done.end())
               ans++;
         }
      }
      return ans;
   }
};
main(){
   Solution ob;
   vector<int> v = {3,1,4,1,5};
   cout << (ob.findPairs(v, 2));
}

入力

{3,1,4,1,5}, 2

出力

2

計算量の目安

このアルゴリズムの時間計算量は O(n log n) です。各要素のセット・マップへの挿入と検索に O(log n) かかり、それを n 個の要素に対して行うためです。空間計算量は O(n) となります。