[백준][1059] 좋은 구간

• Problem Link : https://www.acmicpc.net/problem/1059

문제

• 특정 조건을 만족하는 [A, B]를 좋은 구간이라고 한다.
• 좋은 구간은 쉽게 말해 집합 S의 숫자들이 구간 사이에 위치하지 않아야 좋은 구간이다.

입력

• 첫째 줄, 정수 집합 S의 크기 L이 주어진다.
• 둘째 줄, 집합에 포함된 정수가 주어진다.
• 셋째 줄에는 n이 주어진다.

출력

• n을 포함하는 좋은 구간의 개수를 출력한다.

풀이

위에서 기술했듯이 좋은 구간이란 쉽게 말해 집합 S의 수들이 A, B 구간 사이에 위치하지 않아야 [A, B]를 좋은 구간이라고 할 수 있다. 처음 이 문제를 보고 떠오른 것은 조합으로 풀면 어떨까 생각을 했다.
집합을 정렬 시킨후 우리가 원하는 n이 어떤 구간 사이에 위치해야하는지를 파악하고 어떤 구간인지를 파악한다면 해당 구간의 길이를 구할 수 있다. 그렇다면 그 구간의 수들 중에서 n을 포함한 구간을 뽑는 조합을 구할 수 있지 않을까 했다. n을 포함하는 구간의 조합의 수는 2개씩 뽑은 다음에 n을 벗어나는 구간을 제외시키면 된다.
즉, 2개를 뽑는 경우의 수에서 n의 앞에서만 2개 뽑은 경우, n의 뒤에서만 두 개 뽑는 경우를 제외시키면 된다.

이 때 combination(C)를 다음과 같이 함수로 구현했다.

def combination(l, r):
    return int((r - l + 1) * (r - l) / 2)

이제 공식은 세웠으니 실제 입력을 받고 n이 어느 구간에 포함되는지를 확인하여 [left, right]를 구해서 공식에 대입시킨 결과를 출력하면 된다.

n이 집합 s에 들어있는 경우는 0을 출력해버리고 종료되게 처리를 해줘야한다. 구간을 모두 방문해서 확인하는 시간 낭비를 줄일 수 있다.

코드

import sys

input = sys.stdin.readline

l = int(input())
s = list(map(int, input().split()))
n = int(input())
s.sort()


def combination(l, r):
    return int((r - l + 1) * (r - l) / 2)


if n in s:
    print(0)
else:
    left = 0
    right = 0
    for x in s:
        if x < n:
            left = x
        elif x > n and right == 0:
            right = x
    left += 1
    right -= 1
    print(
        combination(left, right) - combination(left, n - 1) - combination(n + 1, right)
    )