Computer >> 컴퓨터 >  >> 프로그래밍 >> Python

파이썬으로 방정식을 성립시키기 위한 최소 숫자 삽입 횟수 구하기

문제 개요

x+y=z 형태의 방정식을 나타내는 문자열 s가 주어졌다고 가정해 보겠습니다. 이 문제의 목표는 방정식이 실제로 성립하도록 만들기 위해 문자열에 삽입해야 하는 숫자(자릿수)의 최소 개수를 구하는 것입니다.

예를 들어 입력이 s = '2+6=7'이라면 출력은 2가 됩니다. '1'과 '2'를 각각 삽입하여 방정식을 "21+6=27"로 바꾸면 등식이 성립하므로, 필요한 수정 횟수는 총 2회입니다.

해결 접근 방법

이 문제는 세 수 A, B, C의 각 자릿수를 뒤에서부터 한 자리씩 비교하면서 자리올림(carry)을 함께 추적하는 재귀적 동적 계획법(DP)으로 해결할 수 있습니다. 전체 흐름은 다음과 같습니다.

  • 문자열 s를 "+" 기준으로 분할하여 왼쪽 부분은 A에, 나머지는 rest에 저장합니다.
  • rest를 다시 "=" 기준으로 분할하여 왼쪽은 B에, 오른쪽은 C에 저장합니다.
  • dp(len(A)-1, len(B)-1, len(C)-1, 0)을 호출하여 최종 결과를 반환합니다.

dp() 함수의 동작 원리

dp() 함수는 각 문자열의 현재 인덱스 i, j, k와 자리올림 carry를 매개변수로 받아, 가능한 모든 삽입 전략을 탐색하며 최솟값을 계산합니다.

  • 종료 조건: i, j, k가 모두 -1 이하라면 carry가 0일 때 0, 그렇지 않으면 1을 반환합니다.
  • last1, last2, last3: 각 문자열의 현재 자릿수를 의미하며, 인덱스가 0 미만이면 0으로 처리합니다.
  • prefix1, prefix2, prefix3: 각 문자열의 시작부터 현재 위치까지의 정숫값입니다.
  • A와 B가 모두 소진된 경우(i ≤ -1, j ≤ -1): rhs = prefix3 - carry를 계산합니다. rhs ≤ 0이면 |rhs|를, 그렇지 않으면 len(str(rhs))를 반환합니다.
  • C가 소진된 경우(k ≤ -1): len(str(prefix1 + prefix2 + carry))를 반환합니다.

그 외의 경우에는 다음 네 가지 선택지를 모두 고려하여 그중 최솟값을 답으로 사용합니다.

  • 현재 자릿수들의 합이 C의 자릿수와 일치하면 그대로 다음 자리로 진행
  • B 쪽에 숫자를 하나 더 삽입하는 경우
  • A 쪽에 숫자를 하나 더 삽입하는 경우
  • C 쪽에 숫자를 하나 더 삽입하는 경우

구현 예제

다음 코드를 통해 실제 구현을 확인해 보겠습니다.

class Solution:
   def solve(self, s):
      A, rest = s.split("+")
      B, C = rest.split("=")
      def dp(i, j, k, carry):
         if i <= -1 and j <= -1 and k <= -1:
            return 0 if carry == 0 else 1
         last1 = int(A[i]) if i >= 0 else 0
         last2 = int(B[j]) if j >= 0 else 0
         last3 = int(C[k]) if k >= 0 else 0
         prefix1 = int(A[: i + 1]) if i >= 0 else 0
         prefix2 = int(B[: j + 1]) if j >= 0 else 0
         prefix3 = int(C[: k + 1]) if k >= 0 else 0
         if i <= -1 and j <= -1:
            rhs = prefix3 - carry
            if rhs <= 0:
               return abs(rhs)
            if i == -1 or j == -1:
               return len(str(rhs))
            else:
               assert False
         if k <= -1:
            return len(str(prefix1 + prefix2 + carry))
         ans = float("inf")
         carry2, lhs = divmod(carry + last1 + last2, 10)
         if lhs == last3:
            ans = dp(i - 1, j - 1, k - 1, carry2)
         req = last3 - carry - last2
         extra_zeros = max(0, -1 - i)
         carry2 = 1 if req < 0 else 0
         ans = min(ans, 1 + extra_zeros + dp(max(-1, i), j - 1, k - 1, carry2))
         req = last3 - carry - last1
         extra_zeros = max(0, -1 - j)
         carry2 = 1 if req < 0 else 0
         ans = min(ans, 1 + extra_zeros + dp(i - 1, max(-1, j), k - 1, carry2))
         carry2, lhs = divmod(last1 + last2 + carry, 10)
         ans = min(ans, 1 + dp(i - 1, j - 1, k, carry2))
         return ans
      return dp(len(A) - 1, len(B) - 1, len(C) - 1, 0)

ob = Solution()
print (ob.solve('2+6=7'))

입력

'2+6=7'

출력

2

성능 개선 팁

위 구현은 같은 상태(i, j, k, carry)를 여러 번 중복 계산할 수 있습니다. 함수에 @functools.lru_cache(None) 데코레이터를 붙여 메모이제이션을 적용하면, 상태의 개수는 O(len(A) × len(B) × len(C) × 2)로 제한되므로 입력 길이가 길어져도 효율적으로 동작합니다.

마무리

이 알고리즘은 각 자릿수를 오른쪽에서 왼쪽으로 처리하면서, 숫자를 삽입할 수 있는 네 가지 분기를 모두 탐색함으로써 항상 최소 삽입 횟수를 보장합니다. 자리올림 처리와 접두사 값 활용이 핵심 포인트이며, 메모이제이션과 결합하면 실용적인 수준의 성능을 얻을 수 있습니다.