LeetCode 1652. Defuse the Bomb

はじめに

LeetCode 「1652. Defuse the Bomb」の問題をGoで解きました。

問題

https://leetcode.com/problems/defuse-the-bomb/description

問題文

You have a bomb to defuse, and your time is running out!
Your informer will provide you with a circular array code of length of n and a key k.

To decrypt the code, you must replace every number. All the numbers are replaced simultaneously.

If k > 0, replace the ith number with the sum of the next k numbers.
If k < 0, replace the ith number with the sum of the previous k numbers.
If k == 0, replace the ith number with 0.
As code is circular, the next element of code[n-1] is code[0], and the previous element of code[0] is code[n-1].

Given the circular array code and an integer key k, return the decrypted code to defuse the bomb!

和訳
爆弾を解除しなければなりません、時間は迫っています！
情報提供者は、長さ n の循環配列 code とキー k を提供します。

すべての要素を同時に置き換える必要があります。

• k > 0 の場合: i 番目の数値を、次の k 個の要素の合計で置き換える。
• k < 0 の場合: i 番目の数値を、前の k 個の要素の合計で置き換える。
• k == 0 の場合: i 番の数値を 0 に置き換える。

配列 code は循環するので、code[n-1] の次の要素は code[0]、code[0] の前の要素は code[n-1]です。

与えられた円形配列 code と整数 k に基づいて、復号後の配列を返し、爆弾を解除してください！

例

Input: code = [5,7,1,4], k = 3
Output: [12,10,16,13]
Explanation: Each number is replaced by the sum of the next 3 numbers. The decrypted code is [7+1+4, 1+4+5, 4+5+7, 5+7+1]. Notice that the numbers wrap around.

Input: code = [1,2,3,4], k = 0
Output: [0,0,0,0]
Explanation: When k is zero, the numbers are replaced by 0. 

Input: code = [2,4,9,3], k = -2
Output: [12,5,6,13]
Explanation: The decrypted code is [3+9, 2+3, 4+2, 9+4]. Notice that the numbers wrap around again. If k is negative, the sum is of the previous numbers.

制約

• n == code.length
• 1 <= n <= 100
• 1 <= code[i] <= 100
• -(n - 1) <= k <= n - 1

解答1

実装を考える上で、配列が循環することがポイントになりそうです。

例1Input: code = [5,7,1,4], k = 3で考えてみます。
index が 0 の要素は 7(code[1]) + 1(code[2]) + 4(code[3]) で計算できます。
index が 1 の要素は 1(code[2]) + 4(code[3]) + 5(code[0]) で計算したいです。
循環するので、純粋に index を増やした場合 code[4] となるところを code[0] に変換したいです。
この時、4(index) % 4(n)で 0を導くことができます。
同様に code[5] となるところは 5(index) % 4(n)で code[1] に変換できます。

コード

func decrypt(code []int, k int) []int {
    n := len(code)
    decrypted := make([]int, n)

    switch {
    case k > 0:
        for i, _ := range code {
            v := 0
            for j := 1; j <= k; j++ {
                v += code[(i+j)%n]
            }
            decrypted[i] = v
        }
    case k < 0:
        for i, _ := range code {
            v := 0
            for j := -1; j >= k; j-- {
                v += code[(i+j+n) % n]
            }
            decrypted[i] = v
        }
    }

    return decrypted
}

k > 0の時は code[(i+j)%n] で循環配列を考慮したインデックスを計算しています。
k < 0の時はマイナスになることがあるので、+n を追加して code[(i+j+n) % n]でインデックスを計算しています。

なお、newCode := make([]int, n)で各値が0で初期化されているので、k == 0 の分岐は不要です。

計算量

• 時間的計算量：O(n⋅∣k∣)
• 空間的計算量：O(n)

解答2（スライディングウィンドウ）

より効率的なスライディングウィンドウというアプローチがあるので紹介します。

スライディングウィンドウとは

スライディングウィンドウ（Sliding Window）は、配列や文字列を一定の範囲（ウィンドウ）でスキャンしながら、逐次更新する手法です。
全探索を使うとO(n²)になるところを、O(n)に抑えられる場合が多いです。

考え方

1. 最初のウィンドウを作る
   • 今回の問題だと、index 0 に対応する k 個の要素の合計を計算する
2. ウィンドウをずらしながら更新
   • 新しい要素を加え、古い要素を引くことで、効率よく次のウィンドウの値を求める

コード

func decrypt(code []int, k int) []int {
    n := len(code)
    decrypted := make([]int, n)

    if k == 0 {
        return decrypted
    }

    // 初期ウィンドウの合計
    sum := 0
    if k > 0 {
        for i := 1; i <= k; i++ {
            sum += code[i % n]
        }
    } else {
        for i := n + k; i < n; i++ {  // (i = n+k) は (i = n-|k|) と同じ
            sum += code[i]
        }
    }

    // スライディングウィンドウで更新
    for i := 0; i < n; i++ {
        decrypted[i] = sum
        if k > 0 {
            sum -= code[(i+1) % n]         // 古い要素を引く
            sum += code[(i+k+1) % n]       // 新しい要素を加える
        } else {
            sum -= code[(i+k+n) % n]       // 古い要素を引く
            sum += code[i]                 // 新しい要素を加える
        }
    }

    return decrypted
}

まず k の値に応じて初期ウィンドウの合計 sum を求めます。

k > 0 の場合、インデックス 1 から k までの要素を加算します。
k < 0 の場合、末尾側の k 個の要素を加算します。
その後、スライディングウィンドウを適用して各要素を求めます。
ループ内では、sum からウィンドウの古い要素を引き、新しい要素を加えることで次の sum を求めています。

さらに、以下のように改良することもできます。

func decrypt(code []int, k int) []int {
    n := len(code)
    decrypted := make([]int, n)

    if k == 0 {
        return decrypted
    }

    // 初期ウィンドウの合計
    sum := 0
    start, end := 1, k
    if k < 0 {
        start, end = n+k, n-1
    }
    for i := start; i <= end; i++ {
        sum += code[i % n]
    }

    // スライディングウィンドウで更新
    for i := 0; i < n; i++ {
        decrypted[i] = sum
        sum -= code[(i+start) % n]  // 古い要素を引く
        sum += code[(i+end+1) % n]  // 新しい要素を加える
    }

    return decrypted
}

計算量

• 時間的計算量：O(n)
• 空間的計算量：O(n)