CVTE笔试题2：生成zigzag矩阵

其几天参加的CVTE两道笔试题，第二道是生成zigzag矩阵，具体如下
给定数组{1，2，3，4，5，6，7，8，9}

输出
【
1，2，6
3，5，7
4，8，9
】

vector<vector<int>> to_ZigZagMatrix(vector<int>& data)
{
	//参数合法性判断
	int n = data.size();
	if (n == 0)
		throw new exception("输入数组为空");
	if (n == 1)
		return { data };
	int temp = sqrt(n);
	if (temp*temp != n)
		throw new exception("输入数组长度不合法：长度非平方数");
	else
		n = temp;

	//初始化矩阵
	vector<vector<int>>res;
	for (int i = 0; i < n; ++i)
		res.push_back({});
	int count = 0;
	for (int i = 0; i < n*n; ++i)
	{
		if (i%n == 0 && i != 0)
			count++;
		res[count].push_back(0);
	}

	//生成zigzag矩阵
	int rows = n, cols = n;
	int s = 0, t = 0;
	for (int p = 0; p < rows*cols; ++p)
	{
		int row =  p / n;
		int col =  p % n;
		if (row + col <= n - 1)
		{
			s = row + col;
			t = s*(s + 1) / 2;
			if (s % 2 == 0)
				res[row][col] = data[t + col];
			else
				res[row][col] = data[t + row];
		}
		else
		{
			s = n - 1 - row + n - 1 - col;
			t = s*(s + 1) / 2;
			if (s % 2 == 0)
				res[row][col] = data[n*n - 1 - t - (n - 1 - col)];
			else
				res[row][col] = data[n*n - 1 - t - (n - 1 - row)];
		}
	}
	return res;
}