HackerEarth提交Go数组右旋转代码超时 输入处理问题排查

问题现象 #

在HackerEarth平台提交数组右旋转题目的Go语言实现，绝大多数测试用例运行通过，仅1个大规模用例返回*Time limit exceeded（超出时间限制）*报错，需要排查代码尤其是输入处理环节的性能问题，原提交代码如下：

package main

import(
    "fmt"
    "strings"
)

func rotateRight(numbers []int, size int, k int) []int {
    new_numbers := make([]int, size)
    for index, value := range numbers {
        new_numbers[(index + k) % size] = value
    }

    return new_numbers
}

func main() {
    var test_case, size, k int

    fmt.Scanf("%v", &test_case)
    fmt.Scanln()
    for i := 0; i < test_case; i++ {
        fmt.Scanf("%v %v", &size, &k)
        fmt.Scanln()
        
        numbers := make([]int, size)
        for i := 0; i<size; i++ {
            fmt.Scanf("%v", &numbers[i])
        }

        result := rotateRight(numbers, size, k)

        fmt.Println(strings.Trim(fmt.Sprint(result), "[]"))
    }
}

核心性能问题点 #

• 输入输出效率过低：原生fmt包的Scan、Print系列函数默认带严格的格式校验、无缓冲IO，逐元素读取/写入的系统调用开销在数据量达到10^5级别以上时会成为主要耗时，这是本题超时的核心原因。
• 旋转逻辑缺少必要预处理：没有提前对偏移量k做取模运算，当k远大于数组长度时，(index + k) % size会做无意义的大数计算；同时额外开辟了和原数组等长的新切片，内存分配和拷贝的开销在大数据量下也会拖慢运行速度。
• 输出逻辑冗余：通过fmt.Sprint转切片再裁剪方括号的实现，会先生成带格式的完整字符串再做字符串处理，比直接拼接输出内容的开销高很多。

优化后参考实现 #

优化点包括：用bufio做带缓冲的批量IO、提前对k取模、用strings.Builder预拼接输出内容、原地旋转减少内存开销：

package main

import (
	"bufio"
	"fmt"
	"os"
	"strconv"
	"strings"
)

// 原地三次反转实现右旋转，额外空间复杂度O(1)
func rotateRight(nums []int, k int) {
	n := len(nums)
	if n == 0 {
		return
	}
	k = k % n
	if k == 0 {
		return
	}
	reverse(nums, 0, n-1)
	reverse(nums, 0, k-1)
	reverse(nums, k, n-1)
}

func reverse(nums []int, start, end int) {
	for start < end {
		nums[start], nums[end] = nums[end], nums[start]
		start++
		end--
	}
}

func main() {
	// 初始化带缓冲的读写器，批量处理IO
	scanner := bufio.NewScanner(os.Stdin)
	scanner.Split(bufio.ScanWords)
	writer := bufio.NewWriter(os.Stdout)
	defer writer.Flush()

	nextInt := func() int {
		scanner.Scan()
		num, _ := strconv.Atoi(scanner.Text())
		return num
	}

	testCase := nextInt()
	for i := 0; i < testCase; i++ {
		size := nextInt()
		k := nextInt()
		nums := make([]int, size)
		for j := 0; j < size; j++ {
			nums[j] = nextInt()
		}

		rotateRight(nums, k)

		// 预拼接输出字符串，一次性写入
		var builder strings.Builder
		for idx, val := range nums {
			if idx > 0 {
				builder.WriteByte(' ')
			}
			builder.WriteString(strconv.Itoa(val))
		}
		writer.WriteString(builder.String() + "\n")
	}
}

内容的提问来源于stack exchange，提问作者ZamaaN