Skip to content

3267. Count Almost Equal Pairs II

  • Time: $O(n \cdot \log(\max(\texttt{nums}))^4)$
  • Space: $O(\log(\max(\texttt{nums}))^4)$
 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
class Solution {
 public:
  // Similar to 3265. Count Almost Equal Pairs I
  int countPairs(vector<int>& nums) {
    int ans = 0;
    unordered_map<int, int> count;
    const int maxLen = to_string(ranges::max(nums)).length();

    for (const int num : nums) {
      const string digits =
          string(maxLen - to_string(num).length(), '0') + to_string(num);
      for (const int swap : getSwaps(digits))
        ans += count[swap];
      ++count[num];
    }

    return ans;
  }

 private:
  // Returns all possible numbers after 1 or 2 swaps.
  unordered_set<int> getSwaps(const string& digits) {
    const int n = digits.length();
    unordered_set<int> swaps{{stoi(digits)}};

    // Add all numbers after 1 swap.
    for (int i = 0; i < n; ++i)
      for (int j = i + 1; j < n; ++j) {
        string newDigits = digits;
        swap(newDigits[i], newDigits[j]);
        swaps.insert(stoi(newDigits));
      }

    // Add all numbers after 2 swaps.
    for (int i1 = 0; i1 < n; ++i1)
      for (int j1 = i1 + 1; j1 < n; ++j1)
        for (int i2 = 0; i2 < n; ++i2)
          for (int j2 = i2 + 1; j2 < n; ++j2) {
            string newDigits = digits;
            swap(newDigits[i1], newDigits[j1]);
            swap(newDigits[i2], newDigits[j2]);
            swaps.insert(stoi(newDigits));
          }

    return swaps;
  }
};

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
class Solution {
  // Similar to 3265. Count Almost Equal Pairs I
  public int countPairs(int[] nums) {
    int ans = 0;
    Map<Integer, Integer> count = new HashMap<>();
    final int maxLen = String.valueOf(Arrays.stream(nums).max().getAsInt()).length();

    for (final int num : nums) {
      final String digits = String.format("%0" + maxLen + "d", num);
      for (final int swap : getSwaps(digits))
        ans += count.getOrDefault(swap, 0);
      count.merge(num, 1, Integer::sum);
    }

    return ans;
  }

  // Returns all possible numbers after 1 or 2 swaps.
  private Set<Integer> getSwaps(final String digits) {
    final int n = digits.length();
    Set<Integer> swaps = new HashSet<>(Arrays.asList(Integer.parseInt(digits)));

    // Add all numbers after 1 swap.
    for (int i = 0; i < n; ++i)
      for (int j = i + 1; j < n; ++j) {
        char[] newDigits = digits.toCharArray();
        char temp = newDigits[i];
        newDigits[i] = newDigits[j];
        newDigits[j] = temp;
        swaps.add(Integer.parseInt(new String(newDigits)));
      }

    // Add all numbers after 2 swaps.
    for (int i1 = 0; i1 < n; ++i1)
      for (int j1 = i1 + 1; j1 < n; ++j1)
        for (int i2 = 0; i2 < n; ++i2)
          for (int j2 = i2 + 1; j2 < n; ++j2) {
            char[] newDigits = digits.toCharArray();
            char temp = newDigits[i1];
            newDigits[i1] = newDigits[j1];
            newDigits[j1] = temp;
            temp = newDigits[i2];
            newDigits[i2] = newDigits[j2];
            newDigits[j2] = temp;
            swaps.add(Integer.parseInt(new String(newDigits)));
          }

    return swaps;
  }
}

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
class Solution:
  # Similar to 3265. Count Almost Equal Pairs I
  def countPairs(self, nums: list[int]) -> int:
    ans = 0
    count = collections.Counter()
    maxLen = len(str(max(nums)))

    for num in nums:
      digits = list(str(num).zfill(maxLen))
      for swap in self._getSwaps(digits):
        ans += count[swap]
      count[num] += 1

    return ans

  def _getSwaps(self, digits: str) -> set[int]:
    """Returns all possible numbers after 1 or 2 swaps."""
    n = len(digits)
    swaps = set([int(''.join(digits))])

    # Add all numbers after 1 swap.
    for i, j in itertools.combinations(range(n), 2):
      newDigits = digits[:]
      newDigits[i], newDigits[j] = newDigits[j], newDigits[i]
      swaps.add(int(''.join(newDigits)))

    # Add all numbers after 2 swaps.
    for (i1, j1), (i2, j2) in itertools.combinations(
            itertools.combinations(range(n), 2), 2):
      newDigits = digits[:]
      newDigits[i1], newDigits[j1] = newDigits[j1], newDigits[i1]
      newDigits[i2], newDigits[j2] = newDigits[j2], newDigits[i2]
      swaps.add(int(''.join(newDigits)))

    return swaps