[CodeUp/2818] 행운의 뽑기 (python)

Problem : 행운의 뽑기 (codeup.kr)

Status : Solved

Time : 00:45:28

 

---

 

문제 설명

 

뽑기 마니아 유신은 행운의 뽑기를 해보려고 고민 중이다. 뽑기의 규칙은 다음과 같다.

 

A. −10^8이상 10^8이하의 중복되지 않는 숫자가 각각 들어있는 공이 N개 있다.

B. 임의로 서로 다른 공 4개를 뽑은 후, 뽑은 공에 담긴 네 숫자의 합을 구한다.

C. 네 숫자의 합이 정확히 행운의 번호 K와 일치하면 당첨되어 상품을 받는다.

행운의 번호와 공들에 적힌 숫자들은 모두 공개되어 있다.

유신은 최대한 신중하게 소비를 하고 싶기 때문에, 자신이 뽑기를 1회 했을 때 당첨될 확률을 구하고 싶어 한다.

 

입력 및 출력

 

입력

첫 줄에 N, K가 주어진다. (4 ≤ N ≤ 500, −4×10^8 ≤ K ≤ 4×10^8)

두 번째 줄에 각각의 공들에 들어 있는 N개의 숫자가 공백을 기준으로 주어진다.

 

출력

유신이 뽑기에 당첨될 확률을 분수꼴로 출력한다.

단, 반드시 기약분수로 출력해야 하며, 분자가 0일 경우 0을, 분모가 1일 경우 분자만을 출력한다.

 

입력 예시

5 11

5 4 2 1 3

 

출력 예시

1/5

 

---

 

풀이

 

우선, brute-force로 이 문제를 접근해서는 최대 약 500^4개의 경우의 수를 계산해야 하므로, 시간초과가 날 수 밖에 없다. 따라서 다른 풀이를 시도해 보는 게 좋다고 생각했다.

 

이런 문제를 풀 때는 투 포인터 알고리즘을 써 보는게 우선인데... 문제는 4개의 수를 고려해야 한다는 점이다. 그 점 때문에 조금 많이 해맸었는데, 일단 나는 이중 투 포인터 방식대로 접근했다.

우선 투 포인터를 적용하기 위해 뽑기 숫자를 정렬 후, 초기 값으로 제일 왼쪽의 Min값과 제일 오른쪽의 Max값을 잡는다. 남은 두 수는 Min값과 Max값 사이의 리스트에서 내부 투 포인터 알고리즘을 적용해 목표로 하는 K값을 찾게 된다.

 

그리고 그 전에, 다음과 같은 경우를 생각해 볼 수 있다:

 

1. 이 상황에서 최소합은 왼쪽 세 값과 오른쪽 값의 합이 된다. 이 최소합보다 K가 더 작다면, 이 리스트에는 정답이 되는 경우의 수가 없다. 따라서 Max값 포인터를 왼쪽으로 옮겨 주어야 한다.

---

 

2. 또한, 이 상황에서 최대합은 왼쪽 값과 오른쪽 값의 합이 된다. 이 최대합보다 K가 더 크다면, 역시 이 리스트에는 정답이 되는 경우의 수가 없다. 따라서 Min값 포인터를 오른쪽으로 옮겨 주어야 한다.

---

 

3. 1, 2번에 해당되지 않는 경우, 내부 투 포인터 알고리즘을 이용하여 Min, Max값을 포함한 총합이 K가 되는 가짓수를 찾는다. 그 후, 재귀적으로 Min 포인터를 오른쪽으로 옮긴 경우와, Max 포인터를 왼쪽으로 옮긴 경우를 재귀적으로 전부 탐색한다.

---

 

이 때 시간복잡도는 내부 투 포인터 O(N), 외부 투 포인터 최대 O(N^2)이므로 최악의 경우 O(N^3)이다. 중복을 줄이기 위해 visited 배열을 이용해 계산을 마친 경우는 시행하지 않았고, 1, 2경우에 대해 휴리스틱하게 계산 과정을 경감할 수 있어 시간 초과 없이 통과할 수 있었던 것 같다.

 

내부 투 포인터 계산은 우선 리스트를 정확히 둘로 나누고 탐색하는 방식이다. 이 역시 다음과 같은 경우를 생각해 볼 수 있을 것이다.

1. 왼쪽 리스트(최소 리스트)에 두 수가 모두 포함된 경우 : 최소 리스트의 최대 조합은 오른쪽 두 값의 합이므로, 이를 앞선 외부 투 포인터 계산과 같이 생각해본다.  이 값이 목표값보다 크거나 같다면 정답이 존재할 수 있다는 의미이다. 따라서 왼쪽 리스트에 대해 내부 투 포인터 계산을 재귀적으로 시행한다.
2. 오른쪽 리스트(최대 리스트)에 두 수가 모두 포함된 경우 : 앞선 1번 경우와 동일하다.
3. 왼쪽 리스트와 오른쪽 리스트에 두 수가 하나씩 포함된 경우 : 내부 투 포인터 계산을 시행한다.

 

 

그 외에 분모 구하기(조합), 분수 출력을 기약분수로 만들기(최대공약수) 등은 고등 수학 내에서 설명 가능하므로 개략적인 설명은 문제 풀이 코드로 대신한다. 너무 고려하던 게 많아서 코드가 많이 더러워진 기분이 든다....

 

풀이 코드

def gcd(a, b):
  while b > 0:
      a, b = b, a % b
  return a

def two_pointer_search(ball_list, num):
  length = len(ball_list)
  if length < 2 :
    return 0
  if length == 2 :
    return 0 if num != sum(ball_list) else 1

  left_list, right_list = ball_list[:length//2], ball_list[length//2:]
  l_p, r_p = 0, len(right_list)-1
  result = 0

  if len(left_list) >= 2 and sum(left_list[-2:]) >= num :
    result += two_pointer_search(left_list, num)
  if len(right_list) >= 2 and sum(right_list[:2]) <= num :
    result += two_pointer_search(right_list, num)
  
  while l_p < len(left_list) and r_p > -1 :
    new_num = left_list[l_p] + right_list[r_p]
    if new_num == num :
      result += 1
      l_p += 1
      r_p -= 1
    elif new_num < num :
      l_p += 1
    else :
      r_p -= 1
  return result

def total_search(start, end, ball_list, num):
  if start >= end - 2 or visited[start][end]:
    return 0
  visited[start][end] = True
  if ball_list[start] + sum(ball_list[end-2:end+1]) < num :
    return total_search(start+1, end, ball_list, num)
  if sum(ball_list[start:start+3]) + ball_list[end] > num :
    return total_search(start, end-1, ball_list, num)
  result = two_pointer_search(ball_list[start+1:end], K-ball_list[start]-ball_list[end])
  result += total_search(start+1, end, ball_list, num) + total_search(start, end-1, ball_list, num)
  return result

N, K = map(int, input().split())
tot_num = (N * (N-1) * (N-2) * (N-3)) // 24

ball_list = sorted(list(map(int, input().split())))
start, end = 0, len(ball_list)-1
visited = [[False]*N for _ in range(N)]

result = total_search(start, end, ball_list, K)

if result == 0 :
  print(0)
else :
  gcd_num = gcd(result, tot_num)
  if gcd_num == tot_num :
    print(1)
  else :
    print('{:d}/{:d}'.format(result//gcd_num, tot_num//gcd_num))

 

풀이 완료!