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

파이썬으로 연속된 자릿수 차이가 K인 숫자 찾기 — BFS 완전 정리

이번 글에서는 길이가 N인 숫자 중에서 연속된 두 자릿수의 절대 차이가 항상 K가 되는 모든 숫자를 찾는 프로그램을 파이썬으로 구현해 보겠습니다. 단, 정답에 포함되는 숫자는 0 자체를 제외하고 맨 앞자리가 0으로 시작해서는 안 된다는 조건이 있습니다.

예를 들어 입력이 N = 4, K = 7이라면 출력은 다음과 같습니다.

[1818, 2929, 7070, 8181, 9292]

여기서 0707은 실제로 조건을 만족하지만 맨 앞자리가 0으로 시작하기 때문에 유효하지 않은 숫자로 제외됩니다.

문제 해결 접근 방법

이 문제는 너비 우선 탐색(BFS)을 활용하면 깔끔하게 해결할 수 있습니다. 한 자리 숫자부터 시작해서, 매 단계마다 현재 숫자의 마지막 자릿수(lsd)를 기준으로 lsd − K 또는 lsd + K를 새 자릿수로 덧붙인 숫자를 큐에 계속 추가하는 방식입니다. 이렇게 하면 조건을 만족하는 숫자만 골라서 자릿수를 하나씩 늘려갈 수 있습니다.

알고리즘의 동작 과정을 단계별로 정리하면 다음과 같습니다.

  • N이 1이면 한 자리 숫자 전체, 즉 0부터 9까지의 리스트를 그대로 반환합니다.

  • 큐를 생성하고 1부터 9까지의 숫자로 초기화합니다. (0은 앞자리 0 금지 조건 때문에 시작 숫자에서 제외)

  • N − 1번 반복하면서 각 반복마다 다음 작업을 수행합니다.

    • 현재 큐의 크기를 저장합니다.

    • 큐에서 숫자를 하나씩 꺼내 마지막 자릿수 lsd = num mod 10을 계산합니다.

    • lsd − K ≥ 0이면, 새 숫자 num × 10 + (lsd − K)를 큐 뒤에 삽입합니다.

    • K가 0이 아니고 lsd + K ≤ 9이면, 새 숫자 num × 10 + (lsd + K)를 큐 뒤에 삽입합니다.

  • 모든 반복이 끝난 후 큐에 남아 있는 숫자들을 리스트로 변환해 반환합니다.

여기서 K ≠ 0 조건을 확인하는 이유는, K가 0일 때 같은 숫자가 두 번 삽입되어 결과가 중복되는 것을 방지하기 위함입니다.

파이썬 구현 예제

다음은 위 알고리즘을 collections.deque를 사용해 구현한 코드입니다. 데크(deque)를 사용하면 popleft() 연산이 O(1)로 처리되므로 BFS 구현에 적합합니다.

from collections import deque

def solve(N, K):
   if N == 1:
      return list(range(10))
   queue = deque(list(range(1, 10)))
   for n in range(N - 1):
      len_queue = len(queue)
      for j in range(len_queue):
         num = queue.popleft()
         lsd = num % 10
         if lsd - K >= 0:
            queue.append(num * 10 + lsd - K)
         if K and lsd + K <= 9:
            queue.append(num * 10 + lsd + K)
   return list(queue)

N = 4
K = 7
print(solve(N, K))

실행 결과 확인

입력

4, 7

출력

[1818, 2929, 7070, 8181, 9292]

동작 원리 살펴보기

N = 4, K = 7인 경우를 예로 들어 흐름을 살펴보겠습니다. 처음 큐에는 [1, 2, 3, 4, 5, 6, 7, 8, 9]가 들어 있습니다. 첫 번째 반복에서 숫자 1의 마지막 자릿수는 1이므로, 1 − 7은 음수라 제외되고 1 + 7 = 8이므로 18이 큐에 추가됩니다. 이런 식으로 각 단계마다 조건을 만족하는 숫자만 자릿수를 하나씩 늘려가며 확장되고, 세 번째 반복이 끝나면 네 자리 숫자 [1818, 2929, 7070, 8181, 9292]만 남게 됩니다.

이 방식은 가능한 후보만 큐에 유지하므로 불필요한 탐색 없이 정답을 효율적으로 구할 수 있으며, 백트래킹이나 완전 탐색보다 코드도 훨씬 간결합니다.