はじめに

問題の読解は容易です。整数の配列 nums が与えられるので、次の 2 つの条件を満たすペアがいくつあるか数えろ、というものです。

1. i < j
2. |nums[i] - nums[j]| == k

ここで |x| は x の絶対値を表します。

よく見る問題です。計算量を考えなければ、二重ループで当然のごとく解ける問題です。というわけで初手いってみましょう。

初手

二重ループです。ご容赦ください。
https://leetcode.com/problems/count-number-of-pairs-with-absolute-difference-k/submissions/1613654919

解説するまでもありません。問題文の条件に従って二重ループによってカウントし、その数を返すだけです。 Go は math.Abs が用意されていますが、引数が float64 なので取り扱いが面倒です。絶対値を求めるのは簡単なので、毎回 abs 関数を実装するのが良いと思います。どこかからコードスニペットとして拾ってくるももちろんいいと思います。

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19

func countKDifference(nums []int, k int) int {
    ans := 0
    for i := 0; i < len(nums); i++ {
        for j := i + 1; j < len(nums); j++ {
            if abs(nums[i], nums[j])  == k {
                ans += 1
            }
        }
    }
    return ans
}

func abs(m, n int) int {
    v := m - n
    if v < 0 {
        return -v
    }
    return v
}

計算量に何も気を使っていないものの、このコードは Beats 48.65% 🤔です。

最終手

何も誇れたもんじゃないですが、Solutions を見てまま書いた提出です。 すごく簡潔ですね。
https://leetcode.com/problems/count-number-of-pairs-with-absolute-difference-k/submissions/1614667796

1
2
3
4
5
6
7
8
9

func countKDifference(nums []int, k int) int {
    m := make(map[int]int)
    ans := 0
    for _, v := range nums {
        m[v]++
        ans += m[v - k] + m[v + k]
    }
    return ans
}

まずは、絶対値を含む式の変形を発想するところが入口です。つまり、| nums[i] - nums[j] | == k という式を解きます。数1で登場します。 細かくは書きませんが、 nums[i] - nums[j] が正の場合と負の場合が考えられます。負のときは - (nums[i] - nums[j]) == k と表現でき、正のときは単に nums[i] - nums[j] == k と表現できます。 これらは nums[i] + k == nums[j], nums[i] - k == nums[j] と等しいです。 これを利用して今回の問題を解いています。

ループの最初の m[v]++ は、その数字が登場する回数をカウントしています。初期値は 0 なので、一回目は必ず 1 になります。 そして、次が肝であり、先ほどの式の活きるところです。 m[v - k] は nums[i] - k であり、m[v + k] は nums[i] - k です。 絶対値の計算では正と負の両方の場合があるので、m[v - k] は正の場合、m[v + k] は正の場合のその数のカウントです。 たとえば、nums = [1, 2, 2, 1] の1つ目の 2 を取り扱うとき、m[v - k] は m[2 - 1] -> m[1] == 1, m[v + k] は m[2 + 1] -> m[3] == 0、よって 1 です。 これを実際に繰り返していくとそのロジックが理解できると思います。僕はそうして理解しました。というよりそうしないと理解できません。

以上です！