두 배열의 합 - 2143

2143번: 두 배열의 합

문제의 입력은 다음과 같이 주어졌다고 할 때,

$T=5$

$A = {{1, 3, 1, 2}}$

$B = {{1, 3, 2}}$

($A$의 subset의 합) + ($B$의 subset의 합) = $T$가 되도록 하는 경우의 수를 구하는 문제

시행 착오 1 - 메모리 초과 🥲

첫번째로 접근한 방식은

• A에서 가능한 subset의 합들을 A 배열에 추가
• B에서 가능한 subset의 합들을 B 배열에 추가
• (A의 element) + (B의 element) = T가 되는 경우의 수 찾기

T = int(input())
n = int(input())
A = list(map(int, input().split()))
m = int(input())
B = list(map(int, input().split()))

for i in range(n-1):  
  temp = A[i]
  while temp < T:     # 이 부분 문제 발생 -> 마이너스 값 가능성 
    temp += A[i+1]
    A.append(temp)

for i in range(m-1): 
  temp = B[i]
  while temp < T:
    temp += B[i+1]
    B.append(temp)

res = 0
for num in A:        # A 원소 + B 원소 = T가 되는 경우의 수 찾기 
  for num2 in B:     # 이 부분 최적화 필요 -> 이진 탐색 
    if num2 == (T - num):
      res += 1
print(res)

잘 안돌아가서 스터디 팀원 분에게 피드백을 요청했는데 접근 방식은 비슷하게 맞았지만 구현 방식에서 잘못된 부분이 많았다.

맞았습니다! 풀이

• subset 배열 저장하는 부분 수정
• $B = T - A$가 되는 경우 탐색할 때 Binary Search 사용

다음 두 가지를 반영해서 제출하니까 드디어 성공했다!

T = int(input())
n = int(input())
A = list(map(int, input().split()))
m = int(input())
B = list(map(int, input().split()))
subsetA, subsetB = [], []  # subset 합을 저장할 리스트

def binarySearch_lower(target, arr): # lower bound 인덱스 찾기
  left = 0
  right = len(arr) - 1
  while left <= right:
    mid = (left + right) // 2
    if arr[mid] == target:
      if mid == 0 or arr[mid - 1] != target:
        return mid
      else:
        right = mid - 1
    elif arr[mid] > target:
      right = mid - 1
    else:
      left = mid + 1
  return -1

def binarySearch_upper(target, arr): # upper bound 인덱스 찾기
  left = 0
  right = len(arr) - 1
  while left <= right:
    mid = (left + right) // 2
    if arr[mid] == target:
      if mid == len(arr)-1 or arr[mid + 1] != target:
        return mid
      else:
        left = mid + 1
    elif arr[mid] > target:
      right = mid - 1
    else:
      left = mid + 1
  return -1

for i in range(n):  # A subset 저장
  temp = 0
  for j in range(i, n):
    temp += A[j]
    subsetA.append(temp)

for i in range(m):  # B subset 저장
  temp = 0
  for j in range(i, m):
    temp += B[j]
    subsetB.append(temp)

res = 0
subsetB.sort()

for numA in subsetA:  # Binary Search로 경우의 수 찾기
  numB = T - numA
  count = 0
  if binarySearch_lower(numB, subsetB) != -1:
    upper_bound = binarySearch_upper(numB, subsetB)
    lower_bound = binarySearch_lower(numB, subsetB)
    res += (upper_bound - lower_bound + 1)

print(res)

Binary Search로 특정 원소의 개수 구하는 방법

• $lower bound$ (원소가 등장하는 첫번째 인덱스)
• $upper bound$ (원소가 등장하는 마지막 인덱스)를 각각 찾고
• $(upper bound - lower bound) + 1$ = 원소의 갯수

맨 위로 이동하기