1395. Count Number of Teams

[zeikar]

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)

#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