6. Zigzag Conversion

[hroc135]

https://leetcode.com/problems/zigzag-conversion/description/

[oda]

Go で Python の yield 相当のことをするとメモリーリークするという理解で正しいですか?

そうだとすると何か工夫をする必要がありますね。

[hroc135]

あ、channel を Close していないということですね。gemini に聞いたところ、コンテキストを使って関数終了時に goroutine も止めてチャンネルも閉じる方法と、goroutine で必要な回数分 row を生成したときにチャネルを閉じる方法があって、前者の方が複雑な goroutine のライフサイクルの管理に向いているとのことでした。

[hroc135]

たいそうなことになってしまったという印象です。

func convert(s string, numRows int) string {
	if numRows == 1 || len(s) <= numRows {
		return s
	}

    ctx, cancel := context.WithCancel(context.Background())
    defer cancel()
	rowChan := make(chan int)
	go func() {
        defer close(rowChan)
		row := 0
		for {
			for row < numRows-1 {
                select {
                case <-ctx.Done():
                    return
                case rowChan <- row:
                    row++
                }
			}
			for row > 0 {
                select {
                case <-ctx.Done():
                    return
                case rowChan <- row:
                    row--
                }
			}
		}
	}()

	rows := make([][]byte, numRows)
	for i := range len(s) {
		r := <-rowChan
		rows[r] = append(rows[r], s[i])
	}
	var sb strings.Builder
	for _, row := range rows {
		for _, b := range row {
			sb.WriteByte(b)
		}
	}
	return sb.String()
}

[oda]

あ、はい。私も追加で思いつくことが、コンテキストを使うよりは別のチャンネルを作ったほうがシンプルに見えることと、rowChan にサイズ制限をつけることくらいです。select が散るのが嫌ですね。関数化しますか。

rowChan := make(chan int, 100)
done := make(chan struct{})
defer close(done)

[oda]

DeepSeek 製。

func convert(s string, numRows int) string {
    if numRows == 1 || len(s) <= numRows {
        return s
    }

    rowChan := make(chan int)
    done := make(chan struct{})

    go func() {
        generateRows(done, rowChan, numRows)
    }()

    rows := make([][]byte, numRows)
    for i := 0; i < len(s); i++ {
        r := <-rowChan
        rows[r] = append(rows[r], s[i])
    }
    close(done)

    var sb strings.Builder
    for _, row := range rows {
        sb.Write(row)
    }
    return sb.String()
}

func generateRows(done <-chan struct{}, ch chan<- int, numRows int) {
    defer close(ch)
    row := 0
    for {
        for nextRow := row + 1; row < numRows-1; row = nextRow {
            if !sendRow(done, ch, row) { return }
            nextRow = row + 1
        }
        for nextRow := row - 1; row > 0; row = nextRow {
            if !sendRow(done, ch, row) { return }
            nextRow = row - 1
        }
    }
}

func sendRow(done <-chan struct{}, ch chan<- int, row int) bool {
    select {
    case <-done:
        return false
    case ch <- row:
        return true
    }
}

[oda]

大きい方になっているのではないでしょうか。

[hroc135]

列数は
len(s) / (2*numrows-2) * (numrows-1)
で決まるというのを誰かのPRで見ました。
ということは len(s)/2 ですね。

[hroc135]

つまり、step1の場合、
空間計算量: O(numrows * len(s))
時間計算量: O(len(s) * max(len(s), numrows))
あ、「大きい方」ってこういうことですか？

[thonda28]

ここを for i, r := range s {} のように扱わなかったのは s[i] で byte 型としたかったからですかね？

今回の問題設定の場合、文字は英語（大文字・小文字）と , および . で ASCII の範疇なので byte 型で扱っても問題なさそうですが、制約がマルチバイト文字に拡張された場合だとうまく動かなさそうかなと思いました。

rows := make([][]rune, numRows) と宣言して、for _, r := range s {} のように扱えば r が rune 型なのでマルチバイト文字でも適切に扱えそうです。

[thonda28]

isGoingDown: bool を使わずに、1 または -1 をもつ direction: int のようなものを使う書き方だと、上記の switch 文の処理を共通化できます。row += direction とし、折り返し地点が来たら反転させる（direction *= -1 とする）イメージです。

ただ direction は実際には 1 と -1 しか扱わないものの int 型からは2値であることはわからないので、その点が不明瞭なのはデメリットかもしれません。適宜コメントなどで補う必要がありそうです。

[hroc135]

マルチバイト文字への対応も含めて書いてみました。かなりすっきりしたコードになりました。

func convert(s string, numRows int) string {
	if numRows == 1 || len(s) <= numRows {
		return s
	}
	rows := make([][]rune, numRows)
	row := 0
	direction := 1 // direction only contains either 1 or -1
	for _, c := range s {
		rows[row] = append(rows[row], c)
                row += direction
                if row == 0 || row == numRows - 1 {
                        direction *= -1
                }
	}
	var sb strings.Builder
	for r := range rows {
                for _, c := range rows[r] {
                        sb.WriteRune(c)
                }
	}
	return sb.String()
}

[thonda28]

僕の方で想定していたコードとほぼ同じです。また個人的には、処理のブロックごとに適宜改行を入れるのが区切りがわかりやすくて好きです。

今回のコードだと、大きく「エッジケース処理」、「Zigzag に変換する処理」、「文字列を構築する処理」の3つくらいに区切れそうなので、rows := make([][]rune, numRows) と var sb strings.Builder の上辺りに改行を入れるかなと思いました。（あくまで自分だったら、という話なので参考までに。）