2367. Number of Arithmetic Triplets

はじめに

必ず増加していく整数の配列 nums から以下の条件を満たす回数を数えよ、という問題。制限を見ると 3 <= nums.length <= 200 とあったので、これは総当たりでいけるな、とまず思いました。で、Hint には案の定「これは総当たりでいけるかな〜？？」と書いてありました。ええ、いけますとも。

1
2


i < j < k,
nums[j] - nums[i] == diff and nums[k] - nums[j] == diff

というわけで顔を赤くしながら書いたコードが以下。

https://leetcode.com/problems/number-of-arithmetic-triplets/submissions/1623622064

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


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

説明する必要なんて1ミリもございません。条件どおりに総当りしただけです。とりあえず次いきましょう。

2回目

ちょっと考えてみたら答えにたどり着きました。答えと言っても Beats 100% になったもののそんなしっくり来てません。解く前にアルゴリズムが正しいことを確認していないから・・・。初めは O(n) で実装できないか考えたんですが、まずは O(n^2) でも良いからできないか、という風に考えました。その結果、一つ前の状態を管理しながら探索すれば解けることに気づきました。つまり、以下の説明にある入力のとき、nums[1]: 1 に diff が一致するのは nums[2]: 4 です。このとき、前回は 4 であったと記録します。そして、次は 4 に対して diff が一致するインデックスを探します。今回は nums[4]: 7 ですね。この配列は必ず増加することが約束されているので、これ以上探しても diff より大きな差になる値しか見つかりません。よって二重ループで解けました。

1

Input: nums = [0,1,4,6,7,10], diff = 3

https://leetcode.com/problems/number-of-arithmetic-triplets/submissions/1623629925

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


func arithmeticTriplets(nums []int, diff int) int {
    ans := 0
	for i := 0; i < len(nums); i++ {
        prev := nums[i]
		for j := i + 1; j < len(nums); j++ {
            if nums[j] - prev == diff {
                if  prev != nums[i] {
                    ans++
                    break
                }
                prev = nums[j]
            }
		}
	}
    return ans
}

美しい手

いやいやいや・・・。この Solution を見てください。コードがすべてを語っています。かつ、非常にシンプルな考え方です。だからこそ Easy に分類されているのかもしれませんが。いや、単に総当たりで解けるから Easy ですね。計算量という概念が登場しません。

単一ループです。考え方は簡単です。diff * 2 の値が見つかったら、それ以前に n - diff と n - diff * 2 の値がないか確認するだけです。あればカウントアップします。 https://leetcode.com/problems/number-of-arithmetic-triplets/submissions/1623650212

 1
 2
 3
 4
 5
 6
 7
 8
 9
10


func arithmeticTriplets(nums []int, diff int) int {
    cnt, ans := make([]int, 201), 0
    for _, n := range nums {
        if (n >= 2 * diff) {
            ans += cnt[n - diff] & cnt[n - diff * 2]
        }
        cnt[n] = 1
    }
    return ans
}

実際の商用コードではあまり複数の変数の宣言と代入を同時に行うことはありませんが、競技プログラミングならではのシンプルさですね。商用コードはメンテナンスされる前提なので、メンテナンスする人にとって分かりやすいコードにしておくことが望ましいです。その意味では、たとえば以下のように書いたりします

1
2
3
4
5
6
7


var (
    cnt []int
    ans int
)

cnt = make([]int, 201)
ans = 0

無駄に思えるでしょう。しかし「この変数を使います」と「この変数に値を代入します」が分かれていることにより、少し分かりやすくなっています。まあ、この例ではしっくり来ませんね。この説明は他に譲ります。