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

C++로 풀기: 정확히 steps번 이동 후 같은 위치에 머무르는 경우의 수

문제 설명

arrLen 크기의 배열이 하나 있고, 그 배열의 인덱스 0 위치에 포인터가 있다고 가정해 봅시다. 각 단계(step)마다 우리는 배열에서 왼쪽으로 한 칸 이동하거나, 오른쪽으로 한 칸 이동하거나, 현재 위치에 그대로 머무를 수 있습니다.

이제 두 개의 정수 steps와 arrLen이 주어졌을 때, 정확히 steps번의 이동을 마친 후에도 포인터가 인덱스 0에 머물러 있도록 하는 이동 방법의 수를 구해야 합니다. 만약 답이 너무 크다면 10^9 + 7로 나눈 나머지를 반환하면 됩니다.

예를 들어, steps = 3, arrLen = 2라면 출력은 4가 됩니다. 3번의 이동 후 인덱스 0에 머무를 수 있는 서로 다른 방법은 다음 4가지이기 때문입니다.

  • [오른쪽, 왼쪽, 대기]
  • [대기, 오른쪽, 왼쪽]
  • [오른쪽, 대기, 왼쪽]
  • [대기, 대기, 대기]

접근 방법

이 문제는 동적 계획법(DP)과 메모이제이션을 활용하여 해결할 수 있습니다. 핵심 아이디어는 각 단계마다 가능한 모든 위치에 대해 도달할 수 있는 경우의 수를 누적하는 것입니다. 특히 steps/2 + 1보다 멀리 떨어진 위치는 남은 이동 횟수 안에 다시 인덱스 0으로 돌아올 수 없으므로, 탐색 범위를 min(arrLen, steps/2 + 1)로 제한하면 불필요한 연산을 크게 줄일 수 있습니다.

재귀적 풀이의 과정은 다음과 같습니다.

  • m := 10^9 + 7 (모듈러 값)
  • add(a, b) 함수를 정의합니다. 이 함수는 (a mod m + b mod m) mod m을 반환하여 오버플로우를 방지합니다.
  • solve(n, x, pos) 함수를 정의합니다. 여기서 x는 남은 이동 횟수, pos는 현재 위치입니다.
  • x가 0이라면, pos가 0일 때만 true(1)를 반환합니다.
  • dp[pos][n]이 -1이 아니라면(이미 계산된 값이라면), dp[pos][n]을 그대로 반환합니다.
  • ans := 0으로 초기화한 뒤 다음을 수행합니다.
    • pos > 0이면, ans에 solve(n, x - 1, pos - 1)을 더합니다. (왼쪽으로 이동)
    • pos < n - 1이면, ans에 solve(n, x - 1, pos + 1)을 더합니다. (오른쪽으로 이동)
    • 항상 ans에 solve(n, x - 1, pos)를 더합니다. (제자리에 대기)
  • dp[pos][n] := ans를 저장한 후 ans를 반환합니다.

메인 함수에서는 공간 복잡도를 줄이기 위해 2행짜리 DP 배열(롤링 배열 기법)을 사용합니다.

  • x := min(arrLen, steps / 2 + 1)
  • 크기가 2 × (x + 1)인 2차원 배열 dp를 선언하고 0으로 채웁니다.
  • dp[0][0] := 1 (시작 상태 설정)
  • n := arrLen
  • i를 1부터 steps까지 반복하면서:
    • j를 0부터 min(arrLen, steps/2 + 1) 미만까지 반복하면서:
      • x := (i - 1) mod 2, y := i mod 2 (이전 단계의 행과 현재 단계의 행을 구분)
      • dp[y][j] := dp[x][j] (제자리에 머무르는 경우)
      • j - 1 >= 0이면, dp[y][j]에 dp[x][j - 1]을 더합니다. (왼쪽 칸에서 오른쪽으로 이동해 온 경우)
      • j + 1 < n이면, dp[y][j]에 dp[x][j + 1]을 더합니다. (오른쪽 칸에서 왼쪽으로 이동해 온 경우)
  • 최종적으로 dp[steps mod 2][0]을 반환합니다.

예제 코드 (C++)

#include <bits/stdc++.h>
using namespace std;
typedef long long int lli;
const int MOD = 1e9 + 7;
lli add(lli a, lli b){
   return (a % MOD + b % MOD) % MOD;
}
class Solution {
   public:
   vector<vector<int> > dp;
   int solve(int n, int x, int pos = 0){
      if (x == 0) {
         return pos == 0;
      }
      if (dp[pos][n] != -1)
      return dp[pos][n];
      int ans = 0;
      if (pos > 0)
      ans = add(ans, solve(n, x - 1, pos - 1));
      if (pos < n - 1)
      ans = add(ans, solve(n, x - 1, pos + 1));
      ans = add(ans, solve(n, x - 1, pos));
      dp[pos][n] = ans;
      return ans;
   }
   int numWays(int steps, int arrLen){
      int x = min(arrLen, steps / 2 + 1);
      this->dp = vector<vector<int> >(2, vector<int>(x + 1, 0));
      dp[0][0] = 1;
      int n = arrLen;
      for (int i = 1; i <= steps; i++) {
         for (int j = 0; j < min(arrLen, steps / 2 + 1); j++) {
            int x = (i - 1) % 2;
            int y = i % 2;
            dp[y][j] = dp[x][j];
            if (j - 1 >= 0)
            dp[y][j] = add(dp[y][j], dp[x][j - 1]);
            if (j + 1 < n)
            dp[y][j] = add(dp[y][j], dp[x][j + 1]);
         }
      }
      return dp[steps % 2][0];
   }
};
main(){
   Solution ob;
   cout << (ob.numWays(3,2));
}

입력

3, 2

출력

4

정리

이처럼 동적 계획법과 롤링 배열 기법을 함께 활용하면, 시간 복잡도 O(steps × min(arrLen, steps/2 + 1)), 공간 복잡도 O(min(arrLen, steps/2 + 1))로 문제를 효율적으로 해결할 수 있습니다. 탐색 범위를 steps/2 + 1로 제한하는 최적화 덕분에 배열 길이가 매우 크더라도 실제로 고려해야 할 위치의 수는 크게 줄어들게 됩니다.