Computer >> 컴퓨터 >  >> 프로그래밍 >> C++

C++로 배우는 회전 함수(Rotate Function) 문제 풀이: O(n)으로 최댓값 구하기

문제 개요

정수 배열 A와 그 길이 n이 주어진 상황을 가정해 보겠습니다. 배열 A를 시계 방향으로 k칸 회전한 결과를 배열 B(k)라고 할 때, 회전 함수는 다음과 같이 정의됩니다.

F(k) = 0 × B(k)[0] + 1 × B(k)[1] + ... + (n-1) × B(k)[n-1]

목표는 F(0)부터 F(n-1)까지 모든 값 중에서 최댓값을 찾는 것입니다.

예제로 살펴보기

입력이 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

따라서 이 경우 정답은 26입니다.

해결 알고리즘

모든 회전마다 처음부터 다시 계산하면 O(n²)의 시간이 걸려 비효율적입니다. 대신 접두사 합(prefix sum)과 접미사 합(suffix sum)을 활용하면 훨씬 빠르게 답을 구할 수 있습니다.

알고리즘 단계

  1. n을 배열 A의 크기로 설정하고, n이 0이면 0을 반환합니다.
  2. 크기가 n인 두 배열 left(접두사 합)와 right(접미사 합)를 준비합니다.
  3. left[0] = A[0]으로 초기화한 뒤, i를 1부터 n-1까지 순회하며 left[i]에 left[i-1]과 A[i]를 차례로 더해 누적합니다.
  4. right[n-1] = A[n-1]로 초기화한 뒤, i를 n-2부터 0까지 역순으로 순회하며 right[i]에 right[i+1]과 A[i]를 차례로 더해 누적합니다.
  5. rightMul = 0, cnt = n-1로 설정하고, i를 n-1부터 1까지 감소시키며 rightMul에 A[i] × cnt를 더하고 cnt를 1씩 줄입니다.
  6. 크기가 n인 배열 x를 만들고, i를 0부터 n-2까지 순회하며 x[i]에 rightMul을 저장한 후 rightMul에서 right[i+1]을 뺍니다.
  7. leftMul = 0, cnt = 1로 설정하고, i를 0부터 n-2까지 순회하며 leftMul에 A[i] × cnt를 더하고 cnt를 1씩 늘린 뒤, 마지막에 cnt를 하나 감소시킵니다.
  8. i를 n-1부터 1까지 감소시키며 x[i]에 leftMul을 더하고, leftMul에서 A[i-1] × cnt를 뺍니다. 이때 i-2 ≥ 0이면 leftMul에 left[i-2]를 더합니다.
  9. 배열 x의 최댓값을 반환합니다.

C++ 구현 코드

아래는 위 알고리즘을 C++로 구현한 전체 코드입니다.

#include <bits/stdc++.h>
using namespace std;
typedef long long int lli;
class Solution {
    public:
    int maxRotateFunction(vector<int>& A) {
        lli n = A.size();
        if(n == 0) return 0;
        lli ret = 0;
        vector <lli>right(n);
        vector <lli> left(n);
        left[0] = A[0];
        for(lli i = 1; i < n; i++){
            left[i] += left[i - 1];
            left[i] += A[i];
        }
        right[n - 1] = A[n - 1];
        for(lli i = n - 2; i >= 0; i--){
            right[i] += right[i + 1];
            right[i] += A[i];
        }
        lli rightMul = 0;
        lli cnt = n - 1;
        for(lli i = n - 1; i > 0; i--){
            rightMul += (A[i] * cnt);
            cnt--;
        }
        vector <lli> x(n);
        for(lli i = 0; i < n - 1; i++){
            x[i] = rightMul;
            rightMul -= right[i + 1];
        }
        lli leftMul = 0;
        cnt = 1;
        for(lli i = 0; i < n - 1; i++){
            leftMul += A[i] * cnt;
            cnt++;
        }
        cnt--;
        for(lli i = n - 1; i >= 1; i--){
            x[i] += leftMul;
            leftMul -= (A[i - 1] * cnt);
            if(i - 2 >= 0) leftMul += left[i - 2];
        }
        ret = INT_MIN;
        for(lli i = 0; i < x.size(); i++) ret = max(ret, x[i]);
        return ret;
    }
};
main(){
    Solution ob;
    vector<int> v = {4,3,2,6};
    cout <<(ob.maxRotateFunction(v));
}

실행 결과 확인

입력:

[4,3,2,6]

출력:

26

복잡도 분석 및 추가 팁

위 알고리즘은 배열을 상수 번 순회하므로 시간 복잡도는 O(n)이며, 추가 배열을 사용하므로 공간 복잡도 역시 O(n)입니다.

참고로 수학적 성질을 활용하면 코드를 더욱 간결하게 만들 수 있습니다. 배열 전체 합을 S라고 할 때, 인접한 두 회전 값 사이에는 다음 점화식이 성립합니다.

F(k) = F(k-1) + S − n × A[n−k]

이 식을 이용하면 F(0)만 직접 계산한 뒤 나머지 값들을 선형 시간에 유도할 수 있어, 구현이 훨씬 단순해집니다. 오버플로 방지를 위해 long long 타입을 사용하는 점도 주목할 만합니다.