Update count-equal-and-divisible-pairs-in-an-array.cpp

Author: kamyu104

C++/count-equal-and-divisible-pairs-in-an-array.cpp

@@ -1,5 +1,5 @@
 // Time:  O(n * sqrt(k))
-// Space: O(n + sqrt(k))
+// Space: O(n + sqrt(k)), number of factors of k is at most sqrt(k)
 
 // math, number theory
 class Solution {
@@ -26,10 +26,35 @@ class Solution {
     }
 };
 
+// Time:  O(nlogk + n * sqrt(k)^2) = O(n * k)
+// Space: O(n * sqrt(k)), number of factors of k is at most sqrt(k)
+// math, number theory
+class Solution2 {
+public:
+    int countPairs(vector<int>& nums, int k) {
+        unordered_map<int, unordered_map<int, int>> cnts;
+        for (int i = 0; i < size(nums); ++i) {
+            ++cnts[nums[i]][gcd(i, k)];
+        }
+        int result = 0;
+        for (const auto& [_, cnt] : cnts) {
+            for (const auto& [x, c1] : cnt) {
+                for (const auto& [y, c2] : cnt) {
+                    if (x > y || x * y % k) {
+                        continue;
+                    }
+                    result += (x != y) ? c1 * c2 : c1 * (c1 - 1) / 2;
+                }
+            }
+        }
+        return result;
+    }
+};
+
 // Time:  O(n^2)
 // Space: O(n)
 // brute force
-class Solution2 {
+class Solution3 {
 public:
     int countPairs(vector<int>& nums, int k) {
         unordered_map<int, vector<int>> idxs;