LeetCode Rotate Function

LeetCode Rotate Function

Given an array of integers A and let n to be its length.

Assume Bk to be an array obtained by rotating the array A k positions clock-wise, we define a "rotation function" F on A as follow:

F(k) = 0 * Bk[0] + 1 * Bk[1] + ... + (n-1) * Bk[n-1].

Calculate the maximum value of F(0), F(1), ..., F(n-1).

Note:
n is guaranteed to be less than 105.

Example:

A = [4, 3, 2, 6]

F(0) = (0 * 4) + (1 * 3) + (2 * 2) + (3 * 6) = 0 + 3 + 4 + 18 = 25
F(1) = (0 * 6) + (1 * 4) + (2 * 3) + (3 * 2) = 0 + 4 + 6 + 6 = 16
F(2) = (0 * 2) + (1 * 6) + (2 * 4) + (3 * 3) = 0 + 6 + 8 + 9 = 23
F(3) = (0 * 3) + (1 * 2) + (2 * 6) + (3 * 4) = 0 + 2 + 12 + 12 = 26

So the maximum value of F(0), F(1), F(2), F(3) is F(3) = 26.

给定一个数组A,和函数F(k),其中F(k) = 0 * Bk[0] + 1 * Bk[1] + ... + (n-1) * Bk[n-1],数组B为数组A顺时针旋转k次后的新数组。要求所有F(k)的最大值。

仔细观察样例,会发现F(k)=\sum_{i=0}^{n-1}((k+i)\%n)*a_i。于是可以算出所有的F(k),然后求最大值。代码如下:

class Solution {
public:
	int maxRotateFunction(vector<int>& A) {
		if (A.empty())return 0;
		int ans = INT_MIN, n = A.size();
		for (int k = 0; k < n; ++k) {
			int cur = 0;
			for (int i = 0; i < n; ++i) {
				cur += ((k + i) % n)*A[i];
			}
			ans = max(ans, cur);
		}
		return ans;
	}
};

本代码提交TLE,复杂度是O(kn),遇到大数据时超时了。

如果再仔细研究一下样例,会发现一个更优的递推公式:F(k)=F(k-1)+sum-n*A[n-k]。这样直接根据前一项F(k-1),可以在O(1)时间算出下一项F(k)。总的时间复杂度只有O(n)。代码如下:

class Solution {
public:
	int maxRotateFunction(vector<int>& A) {
		if (A.empty())return 0;
		int n = A.size(), sum = 0, f0 = 0;
		for (int i = 0; i < n; ++i) {
			sum += A[i];
			f0 += i*A[i];
		}
		int ans = f0, pre = f0;
		for (int i = 1; i < n; ++i) {
			int cur = pre + sum - n * A[n - i];
			ans = max(ans, cur);
			pre = cur;
		}
		return ans;
	}
};

本代码提交AC,用时16MS。

Leave a Reply

Your email address will not be published. Required fields are marked *