1년 전 즈음에 도전했던 문제인데, 풀지 못한 기억이 있어서 다시 도전
import sys
input = sys.stdin.readline
def main():
N = int(input())
arr = [int(input()) for _ in range(N)]
# 아직 오른쪽을 볼 가능성이 있는 것들의 모음
stack = [float("inf")]
total = 0
for x in arr:
# 아직 stack에 있는 것들은 오른쪽을 볼 수 있음
if stack[-1] >= x:
stack.append(x)
# stack에 있는 것들 중 오른쪽을 볼 수 없는 것들 추리기
else:
brr = []
while stack[-1] < x:
brr.append(stack.pop())
# brr에 들어있는 것들과 x는 서로 볼 수 있음
total += len(brr)
# brr에 들어간 것끼리 서로 볼 수 있는지 확인
# 연속된 같은 수들을 따로 처리할 필요가 있어서 이러한 과정이 필요
crr = [brr[0]]
cnt = [1]
# brr = [100, 100, 100, 101, 101] => [100, 101], [3,2]
for i in range(1, len(brr)):
if crr[-1] == brr[i]:
cnt[-1] += 1
else:
crr.append(brr[i])
cnt.append(1)
for i in range(len(crr)):
# 서로 같은 수는 서로 마주 볼 수 있음
total += (cnt[i] * (cnt[i] - 1)) // 2
# 바로 왼쪽의 높은 수와는 마주볼 수 있음
# 하나와만 마주몰 수 있으므로 += cnt[i]
# 제일 왼쪽은 그 왼쪽의 수들과 모두 count가 끝났을 때는 + 하지 않음
# 즉, stack에 수가 남아있지 않으면서 자기가 brr 제일 왼쪽일 때는 + 하지 않음
if len(stack) > 1:
total += cnt[i]
else:
if i != (len(crr) - 1):
total += cnt[i]
stack.append(x)
# 작업 끝나고 stack에 남아있는 것들에 대해서도 한번 더 서로 바라볼 수 있는지 확인하는 작업을 수행
# 단 맨앞의 무한대는 제외
# 위의 brr, crr의 코드를 가져왔는데 brr은 stack에서 pop하면서 저장해서, stack의 역순이 되기 때문에 reverse를 한번 해줘야 함
# 그래야 맨 마지막에 stack에서 제일 왼쪽의 숫자가 들어가서 같은 동작을 함 (이 부분에서 실수를 해서 틀렸었음)
stack.reverse()
brr = stack
crr = [brr[0]]
cnt = [1]
for i in range(1, len(brr) - 1):
if crr[-1] == brr[i]:
cnt[-1] += 1
else:
crr.append(brr[i])
cnt.append(1)
for i in range(len(crr)):
# 서로 같은 수는 서로 마주 볼 수 있음
total += (cnt[i] * (cnt[i] - 1)) // 2
# 바로 왼쪽의 높은 수와는 마주볼 수 있음
# 하나와만 마주몰 수 있으므로 += cnt[i]
# 제일 왼쪽은 그 왼쪽에 수들과의 count가 끝났으므로 + 하지 않음
if i != (len(crr) - 1):
total += cnt[i]
print(total)
main()
위 코드가 처음 통과한 코드
우선, 배열을 새로 만들어 중복된 숫자의 개수를 계산하고 활용하는 과정이 번거로웠음
stack에서 빠진 수들에 대해서 서로 간의 쌍의 개수를 계산을 다시 하는데, 거기서 또 조건문이 들어가는게 깔끔하지 못하다 생각하여 코드 개선
아래 코드가 개선된 코드
위 코드에서 번거로웠던 중복된 숫자의 개선을 새는 것을, 이중 배열을 이용하여 stack 안에서 하도록 개선
그리고 모든 개수 계산을 일관성있게 stack 에서 넣고 뺄 때 하도록 하여 코드의 길이 및 가독성 개선
import sys
input = sys.stdin.readline
def main():
N = int(input())
arr = [int(input()) for _ in range(N)]
total = 0
stack = [[arr[0], 1]]
for i in range(1, len(arr)):
cur = arr[i]
empty = False
# 더 이상 오른쪽을 볼 수 없는 것들은 stack에서 제외
while stack[-1][0] < cur:
# cur와는 서로 마주 볼 수 있음
total += stack[-1][1]
stack.pop()
# 중간에 stack이 비면
if len(stack) == 0:
empty = True
stack.append([cur, 1])
if empty:
continue
if stack[-1][0] == cur:
# 같은 것끼리는 서로 볼 수 있으니 stack[-1][1]을 total에 더해줌
total += stack[-1][1]
# 왼쪽에 더 큰게 있으면 마주 볼 수 있으므로 + 1
if len(stack) > 1:
total += 1
stack[-1][1] += 1
else:
# 마지막 하나와만 서로 마주볼 수 있음
total += 1
stack.append([cur, 1])
print(total)
main()