2200. Find All K-Distant Indices in an Array

はじめに

少し文章読解の難易度が他の Easy より高い問題です。とはいえ AtCoder に慣れた諸氏であれば読解の難しさはないかもしれません。何せ、単に与えられた条件のとおりのコードを書けば PASS するので。
https://leetcode.com/problems/find-all-k-distant-indices-in-an-array/description/

|i - j| <= k and nums[j] == key という条件を満たすインデックス i を整数の配列 nums から列挙せよ、という問題です。つまり、配列すべてのインデックスがその条件を満たすなら、すべてのインデックスを含む配列 [0, 1, 2,..., len(nums)-1] を返せば OK です。

初手

とりあえず素直な実装です。ナイーブな実装とも言いますが、ナイーブと書くと何かを許されるような錯覚を覚えるので、素直な実装と呼びます。
https://leetcode.com/problems/find-all-k-distant-indices-in-an-array/submissions/1615723359

 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


func findKDistantIndices(nums []int, key int, k int) []int {
    ans := []int{}
    for i := 0; i < len(nums); i++ {
        for j := 0; j < len(nums); j++ {
            if check(i, j, nums[j], key, k) {
                // fmt.Println("true")
                ans = append(ans, i)
                break
            }
        }
    }
    return ans
}

func check(i, j, n, key, k int) bool {
    // fmt.Println(i, j, n, key)
    return abs(i, j) <= k && n == key 
}

func abs(i, j int) int {
    if i - j < 0 {
        return -(i - j)
    }
    return i - j
}

テストケースの nums の要素数が増えたとき、デバッグとして書いているプリント文によってエラーが出たため、コメントアウトした状態で提出しています。単に、二重ループの中で条件に合致するインデックスがあればスライスに追加する、それだけのコードです。注意点として、i == j の場合も OK なのと、append 後は break することです。 i == j が OK というのがしっくり来ない人がいるかもしれませんが、問題文に書かれている条件に i ≠ j はありませんので、当然 OK なのです。

ただ、当然のごとく O(n^2) の計算量となっており、改善の余地があります。生成 AI は「頻度マップ」と呼んでいましたが頻度カウンターを使った実装により改善ができるのかもしれません。後日書きます。

二回目

実は二回目を出す前に Solutions を見て理想の形での提出はしています。しかし、一夜明けて自分でも書けないと意味がないことから、何も見ずに書いたコードがこちらです。

weak beats

https://leetcode.com/problems/find-all-k-distant-indices-in-an-array/submissions/1620861465

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


func findKDistantIndices(nums []int, key int, k int) []int {
    ans := make(map[int]bool)
    for i, v := range nums {
        if v == key {
            start := max(0, i - k)
            end := min(len(nums) - 1, i + k)
            for j := start; j <= end; j++ {
                ans[j] = true
            }
        }
    }
    ans2 := []int{}
    for k, _ := range ans {
        ans2 = append(ans2, k)
    }
    sort.Ints(ans2)
    return ans2
    
}

考え方の基礎はこの Solution を参考にしています。

nums[j] == key という条件から、この問題は key に一致するインデックスを起点とした探索をすればよいと考えることができます。1つ目の条件は |i - j| <= k であって、これは j <= i - k と j <= i + k の二通りが考えられます（絶対値記号の外し方より）i - k と i + k は、i はインデックスを指すので、そのインデックスの左右 k 分の距離と理解できます。よって、やるべきことは key が見つかったとき、その左右 k 分の幅を記録していく、という作業です。

上記のコードで言うと、start, end が探索の始点と終点です。重複を考慮すればよかったんですが、上記はとりあえず map に詰め込んで最後に重複なく取り出す、という流れです。start, end の間の区間に見つかったインデックスを map に入れておくことで、最後にキーをすべて取り出せば答えになります。計算量は、LeetCode AI によれば O(Nk+NLogN) でした。遅い。大きく負けている。

not enough beats

最終手

結局最後まで Solutions にお世話になってしまいました。重複をなくす処理の max が思いつかず。道筋立てて考えると当然なんですが、自分では解いておらず、Solutions を見て回答したからこうなってますね。つまり理解できていない。

一つ前のコードとの差異は、答えの変数を map から整数のスライスに変えたことと重複をなくす処理を加えたことです。これで Beats 100% です。不甲斐ない気持ちでいっぱい。

https://leetcode.com/problems/find-all-k-distant-indices-in-an-array/submissions/1621162468

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


func findKDistantIndices(nums []int, key int, k int) []int {
	ans := []int{}
	for i, v := range nums {
		if v == key {
			start := max(0, i-k)
			// 重複をなくす
			if len(ans) > 0 {
				start = max(ans[len(ans)-1]+1, start)
			}
			end := min(len(nums)-1, i+k)
			for j := start; j <= end; j++ {
				ans = append(ans, j)
			}
		}
	}
	return ans

}

extra 手

他の Solution を見てきれいだなと思い真似しました。 TC（Time Complexity）は O(n) です。from, to を1行に書いたのと、現在どこかを記録する prev 変数の導入によりコードがシンプルになっています。やはり Simplicity is the ultimate sophistication ですね。N.Y. の MoMA に着ていった T シャツにプリントされた文字です。学芸員さんに「良いメッセージだね」って言われたのが今でも記憶に残っています。はい。

https://leetcode.com/problems/find-all-k-distant-indices-in-an-array/submissions/1622022845

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


func findKDistantIndices(nums []int, key int, k int) []int {
	ans := []int{}
	prev := -1
	for i, v := range nums {
		if v == key {
			from, to := max(i - k, 0, prev + 1), min(i + k, len(nums) - 1)
            for ; from <= to; from++ {
                ans = append(ans, from)
                prev = from
            }
		}
	}
    return ans
}