Problem link
https://leetcode.com/problems/count-number-of-teams/
Problem Summary
연속으로 작아지거나 커지는 3 크기의 부분 배열을 구하는 문제.
Solution
브루트 포스는 O(n^3)으로 불가. 약간 스마트한 브루트 포스로 풀린다. O(n^2)
배열의 크기가 3이므로 중간을 기준으로 왼쪽은 더 작은 값, 오른쪽은 더 큰 값을 고르면 된다. 비슷하게 중간을 기준으로 왼쪽은 더 큰 값, 오른쪽은 더 작은 값을 갖는 배열의 개수를 구해서 더하면 정답이 된다.
왼쪽은 더 작은 값의 개수와 오른쪽은 더 큰 값의 개수를 곱하면 전체 개수가 된다.
개선된 풀이
O(n log m) 풀이가 있는데 펜윅 트리 (BIT)를 쓰면 된다. (m은 rating의 최대로 10^4)
https://github.com/zeikar/leetcode/issues/12 이문제와 거의 비슷하다.
왼쪽을 기준으로 설명하자면 수가 하나 나올 때마다 그 값에 해당하는 인덱스를 1 증가시킨다. 중간을 기준으로 더 작은 값들의 개수를 구하려면 0부터 중간값-1 까지 전체 합을 구하면 되고 이는 펜윅 트리로 logm만에 구해진다.
다만 실제 면접이나 코딩테스트에서 펜윅 트리를 구현하라고 하면 쉽진 않을것 같다..
Source Code
class Solution:
def numTeams(self, rating: List[int]) -> int:
n = len(rating)
ans = 0
for i in range(1, n-1):
lessleft = 0
moreleft = 0
for l in range(0, i):
if rating[l] < rating[i]:
lessleft += 1
else:
moreleft += 1
lessright = 0
moreright = 0
for r in range(i + 1, n):
if rating[r] < rating[i]:
lessright += 1
else:
moreright += 1
ans += lessleft * moreright
ans += moreleft * lessright
return ans