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

파이썬으로 합이 k로 나누어 떨어지는 자연수 쌍의 개수 구하기

숫자 n과 값 k가 주어졌을 때, 1부터 N까지의 자연수로 이루어진 배열 A에서 i < j를 만족하는 두 원소 A[i]와 A[j]의 합이 k로 나누어 떨어지는 쌍의 총 개수를 구하는 문제를 생각해 봅시다.

예를 들어, n = 10, k = 4가 입력으로 주어지면 출력은 10이 됩니다. 합이 4로 나누어 떨어지는 쌍이 정확히 10개 존재하기 때문입니다.

해당 쌍들은 다음과 같습니다: (1,3), (1,7), (2,6), (2,10), (3,5), (3,9), (4,8), (5,7), (6,10), (7,9)

해결 접근 방법

이 문제를 효율적으로 해결하기 위해 나머지(rest)의 분포를 이용하는 방법을 사용합니다. 두 수의 합이 k로 나누어 떨어지려면, 두 수를 k로 나눈 나머지의 합이 0 또는 k가 되어야 합니다. 이를 활용한 단계별 풀이는 다음과 같습니다.

  • 1단계: m := (n / k)의 몫, r := n mod k를 계산합니다.
  • 2단계: 빈 맵(딕셔너리) b를 생성합니다.
  • 3단계: 0부터 k-1까지의 각 나머지 i에 대해 b[i] := m으로 초기화합니다. 이는 각 나머지 값이 최소 m번씩 등장한다는 의미입니다.
  • 4단계: m*k+1부터 n까지의 수 i에 대해 나머지 j := i mod k를 구하고, b[j]를 1씩 증가시켜 나머지 분포를 보정합니다.
  • 5단계: 카운터 c := 0으로 초기화합니다.
  • 6단계: 0부터 k-1까지 각 i에 대해 다음을 수행합니다.
    • i1 := i, i2 := (k - i) mod k로 설정합니다.
    • i1과 i2가 같으면(즉, 나머지가 0이거나 k/2인 경우), 같은 그룹 내에서 두 개를 뽑는 경우의 수인 b[i] * (b[i]-1)을 더합니다.
    • 그렇지 않으면 서로 다른 그룹에서 뽑는 경우의 수인 b[i1] * b[i2]를 더합니다.
  • 7단계: 각 쌍이 두 번씩 계산되었으므로 c / 2의 몫을 반환합니다.

예제 코드

아래 파이썬 구현을 통해 더 잘 이해해 보겠습니다.

def solve(n, k):
    m = n // k
    r = n % k

    b = {}
    for i in range(k):
        b[i] = m
    for i in range(m*k+1, n+1):
        j = i % k
        b[j] = b[j] + 1

    c = 0
    for i in range(k):
        i1 = i
        i2 = (k - i) % k
        if i1 == i2:
            c = c + b[i] * (b[i]-1)
        else:
            c = c + b[i1] * (b[i2])
    return c//2

n = 10
k = 4
print(solve(n, k))

입력

n = 10, k = 4

출력

10

시간 복잡도

이 알고리즘의 시간 복잡도는 O(n + k)입니다. 나머지 분포를 계산하는 데 O(n)이 소요되고, 쌍의 개수를 세는 데 O(k)가 소요되기 때문입니다. 모든 쌍을 직접 확인하는 O(n²) 브루트 포스 방식보다 훨씬 효율적입니다.