package railfence
import (
	"strings"
)
func Encode(input string, rails int) string {
	rows := make([]string, rails)
	var flip bool
	var row int
	for _, ch := range input {
		rows[row] += string(ch)
		if flip {
			row--
		} else {
			row++
		}
		if row == 0 || row == rails-1 {
			flip = !flip
		}
	}
	var ret strings.Builder
	for _, line := range rows {
		ret.WriteString(line)
	}
	return ret.String()
}
func Decode(input string, rails int) string {
	// detect row lengths
	length := make([]int, rails)
	var flip bool
	var row int
	for i := 0; i < len(input); i++ {
		length[row]++
		if flip {
			row--
		} else {
			row++
		}
		if row == 0 || row == rails-1 {
			flip = !flip
		}
	}
	// fill the rows
	rows := make([]string, rails)
	for input, i := input, 0; i < rails; i++ {
		rows[i] = input[:length[i]]
		input = input[length[i]:]
	}
	// zig-zag
	var ret strings.Builder
	flip = false
	row = 0
	for i := 0; i < len(input); i++ {
		ret.WriteByte(rows[row][0])
		rows[row] = rows[row][1:]
		if flip {
			row--
		} else {
			row++
		}
		if row == 0 || row == rails-1 {
			flip = !flip
		}
	}
	return ret.String()
}