はじめに

整数の配列 nums と整数 target が与えられるので、0 <= i < j <= n かつ nums[i] + nums[j] < target を満たす組の数をカウントせよ、という問題。 簡単そうですよね。簡単です。 とりあえず、いわゆる brute force で解けます。

初手

解いちゃうのどうなの？と思いますが、とりあえず解く方法が分かっていたので二重ループで解きました。時間計算量は最良ではないです。

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11

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

i を固定して j をカウントアップし、一つ一つ条件に合う組をカウントしていく方法です。直感的です。

最良と思われる手

Two Pointers を使います。 そして、0 <= i < j <= n から勝手に課してしまった制約を取り払います。 それは「与えられた nums の配列は並び替えずそのまま使う」というものです。 この制約は、あくまで 2 つのインデックスの大小関係のみを定義しており、与えられた配列 nums の順番に制限を課すものではありません。 つまり、並び替えてもいいのです。ソートしましょう。ソートすれば、Two Pointers で解けます。

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15

func countPairs(nums []int, target int) int {
    slices.Sort(nums)
    ans := 0
    left := 0
    right := len(nums) - 1
    for left < right {
        if nums[left]+nums[right] < target {
            ans += right - left
            left++
        } else {
            right--
        }
    }
    return ans
}

正直なところ、解き方は以下に詳しく、私もこれを見たので、これ以上の分かりやすい説明はできません。

https://leetcode.com/problems/count-pairs-whose-sum-is-less-than-target/solutions/3933451/two-pointers-approach-easy-to-understand-in-9-languages

が、あえて私の言葉で説明します。

とりあえず、昇順にソートします。話はそれからです。 昇順にソートすると、一番左と一番右は、もちろんそれぞれ最小値と最大値となります。 あとは、個別のケースを考えます。まず、たとえば、最小値 ＋ 最大値 < target となった場合、どうでしょう？ 最小値の方のポインタを右にずらした場合、一つ前と同じかそれより大きい合計になることが明らかです。 なので、最大値を指す右側のポインタを一つ左にずらします。すると、最大値と同じかそれより小さい値が指し示されるので、target より合計が小さくなる可能性があります。

左のポインタを left、右のポインタを right としたとき、nums[left] + nums[right] < target となった場合はどうすべきでしょう？ この場合、right をいくら左に移動していっても、必ず target より小さくなります。これは、配列が昇順にソートされているからです。 なので、right - left 分だけ条件を満たす組があることが分かります。つまり、ans += right - left となります。

これで左側と右側両方のポインタの動きが理解できたので、後はそれをコードに起こすだけです。

この方法は、最初のソートによる O(nlogn) と、要素分だけポインタを動かすので O(n) の時間計算量となりますが、O(nlogn) としてしまって良いと思います。 初手は二重ループでしたので O(n^2) でした。十分高速ですね。