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

Python으로 선형 방정식 ax + by = n의 해가 존재하는 계수 쌍 (a, b) 개수 구하기

문제 개요

값 n이 주어졌을 때, 방정식 a·x + b·y = n이 적어도 하나의 해를 갖도록 하는 쌍 (a, b)(단, a < b)의 개수를 구하는 것이 목표입니다.

예를 들어 n = 4라면 정답은 2입니다. 조건을 만족하는 유효한 쌍은 (1, 2)와 (1, 3)뿐이기 때문입니다.

풀이 접근 방법

이 문제는 각 수의 약수를 미리 계산해 두고, 작은 계수 a부터 차례대로 탐색하면서 중복을 제거하는 방식으로 효율적으로 해결할 수 있습니다. 단계별 과정은 다음과 같습니다.

  • 함수 divisors_gen()을 정의합니다. 이 함수는 n을 인자로 받습니다.
  • divs := 크기가 n+1인 리스트의 리스트를 만들고, 각 내부 리스트는 1로 초기화합니다.
  • divs[0] := 요소가 0 하나뿐인 리스트로 설정합니다.
  • i를 2부터 n까지 반복하면서:
    • j를 1부터 (n // i) + 1 범위까지 반복하며, divs[i * j] 리스트의 끝에 i를 추가합니다.
  • 모든 내부 리스트를 역순으로 뒤집어 divs를 반환합니다. (각 수의 약수가 내림차순으로 정렬됩니다.)

메인 함수(solve)에서는 다음을 수행합니다.

  • result := 0으로 초기화합니다.
  • d_cache := divisors_gen(n+1)을 호출해 약수 테이블을 미리 생성합니다.
  • a를 1부터 n-1까지 반복하면서:
    • i := 1로 초기화하고, s := 빈 집합(set)을 만듭니다.
    • a * i < n인 동안 다음을 반복합니다.
      • b := n - a * i 로 나머지 값을 구합니다.
      • d_cache[b]에 저장된 각 약수 d에 대해:
        • d > a이면, d가 아직 집합 s에 없을 때 result를 1 증가시킵니다. (중복 카운트 방지)
        • d ≤ a이면 내부 반복문을 즉시 종료합니다. (약수가 내림차순 정렬되어 있으므로 이후는 확인할 필요가 없습니다.)
        • d를 집합 s에 추가합니다.
      • i를 1 증가시킵니다.
  • 최종 result를 반환합니다.

동작 원리

핵심 아이디어는 다음과 같습니다. a·x + b·y = n에서 y = 1, x = i라고 가정하면 b = n − a·i가 됩니다. 이때 b의 약수 d(d > a)를 선택하면 b = d · (b/d)이므로 a·i + d·(b/d) = n이 성립합니다. 즉, 쌍 (a, d) 역시 유효한 해를 갖는 계수 쌍이 됩니다. 집합 s를 활용하면 서로 다른 i 값에서 같은 쌍이 여러 번 세어지는 것을 방지할 수 있습니다.

예제 코드

아래 구현을 통해 더 자세히 이해해 보겠습니다.

def divisors_gen(n):
   divs = [[1] for x in range(0, n + 1)]
   divs[0] = [0]
   for i in range(2, n + 1):
      for j in range(1, n // i + 1):
         divs[i * j].append(i)
   return [i[::-1] for i in divs]

def solve(n):
   result = 0
   d_cache = divisors_gen(n+1)

   for a in range(1, n):
      i = 1
      s = set([])
      while a*i < n:
         b = n - a*i
         for d in d_cache[b]:
            if d > a:
               if d not in s:
                  result += 1
            else:
               break
            s.add(d)
         i += 1
   return result

n = 4
print(solve(n))

입력

4

출력

2