[PS] BOJ 7774 / 콘센트

[PS] BOJ 7774 / 콘센트
문제 링크: https://www.acmicpc.net/problem/7774
Thumbnail: Photo by Ryan Quintal (Unsplash)

콘센트 규격이 두가지가 섞여 있는 동네라니... 왜 하나로 통일하지 않은 걸까요?

생각한 과정

처음에 A 단자 1개만 있고, A->B 멀티탭과 B->A 멀티탭을 받아 최대한 많은 A 단자를 만들어야 합니다.

A 단자를 최대한 많이 만들기 위해서는, B 단자는 가능한 많이 사용하되 A->B->A로 돌아올 수 없는 경우는 굳이 A단자에 A->B 케이블을 연결할 필요가 없음을 생각하면 됩니다.

이때, 가장 많은 단자를 얻기 위해서는 케이블을 단자가 많은 것부터 써야 합니다.
-> 정렬 필요!

풀이

아래 코드에서는 가독성을 위해 N을 A2B_CNT, M을 B2A_CNT로 사용했습니다.

먼저, A->B 케이블과 B->A 케이블은 단자 수가 많은 순으로 꺼내 쓸 수 있도록, 내림차순 정렬해 둡니다.

이후, A->B 케이블과 B->A 케이블이 모두 남아있는 동안 아래의 과정을 반복합니다:

  • B 단자가 없다면
    • A->B 멀티탭을 1개 연결합니다.
  • B단자가 남아있고, B->A 케이블이 남아 있는 동안
    • B->A 멀티탭을 계속 연결합니다.

전체 코드

from sys import stdin

A2B_CNT, B2A_CNT = map(int, stdin.readline().split())
A2B = sorted(map(int, stdin.readline().split()), reverse=True)
B2A = sorted(map(int, stdin.readline().split()), reverse=True)

A2B_IDX = 0
B2A_IDX = 0
a_cnt = 1
b_cnt = 0
while A2B_IDX < A2B_CNT and B2A_IDX < B2A_CNT:      # 양쪽 모두 여분의 멀티탭이 있는 동안
    if b_cnt == 0:                                  # 남은 B 단자가 없는 경우
        b_cnt += A2B[A2B_IDX]                       # 가장 많은 수의 B단자를 가진 A->B 멀티탭을 1개 사용
        A2B_IDX += 1
        a_cnt -= 1
    while B2A_IDX < B2A_CNT and b_cnt > 0:          # B단자가 남아있고, B->A 멀티탭이 남아있는 동안
        a_cnt += B2A[B2A_IDX]                       # 가장 많은 수의 A단자를 가진 B->A 멀티탭을 가능한만큼 사용
        B2A_IDX += 1                                # 남은 B단자에 전부 연결하거나, 남아있는 B->A 멀티탭을 전부 쓸 때까지 반복.
        b_cnt -= 1

print(a_cnt)

solution.py

개선할 수 있는 부분?

B 단자에 B->A 케이블을 연결하는 부분은, 굳이 반복하지 않고도 계산할 수 있습니다.

while B2A_IDX < B2A_CNT and b_cnt > 0:
        a_cnt += B2A[B2A_IDX]
        B2A_IDX += 1
        b_cnt -= 1

before

diff = b_left if b_cnt > (b_left := B2A_CNT - B2A_IDX) else b_cnt
a_cnt += sum(B2A[B2A_IDX:B2A_IDX + diff])
B2A_IDX += diff
b_cnt = 0

after