• 공식 블로그에 공식 해설이 이미 있기는 하다.
• 그러나 해설에서는 세그먼트 트리를 이용해서 O(nlogn)에 푸는 방법을 설명하고 있지만, 수식을 잘 정리하면 O(n)에 푸는 것도 가능하다.
  • 사실 수식을 O(n)형태로 정리하는데 생각하는 시간이 상당히 오래 걸렸다. 실제로 대회같은 상황에서 이런 문제가 나오면 빠르게 세그먼트 트리 아이디어를 떠올려서 구현하는게 효과적일 것이고, 연습할때도 그런 방향으로 연습을 하는게 더 도움이 될 것같다.
• 기본 아이디어는, 문자열 s의 아름다움은 s의 길이에서, s의 첫글자와 마지막 글자가 같으면 1을 빼고, s의 첫 두글자와 마지막 두글자가 같으면 1을 또 빼고, … 이런 방식으로 계산하는 것이다.
  • beauty(s) = len(s) - (1 if s[0] == s[-1] else 0) - (1 if s[:1] == s[-2:] else 0) - …
• f[x]를 첫 x개 글자와 마지막 x개 글자가 같은 부분문자열의 갯수라고 하면, 모든 부분문자열 si의 아름다움의 합 sum(beauty(si) = sum(len(si)) - f[1] - f[2] - …. 이 된다.
• $\sum_i(len(s_i)) = \sum_{i=1}^n{\sum_{j=i+1}^n(j - i)} = \frac{n(n-1)(n+1)}{6} $ 은 쉽게 유도 가능
• f[x] = fa[x]+fb[x]+.. 이런식으로 문자별로 쪼개서 생각해보자. fa[x]는 첫 x개 글자와 마지막 x개 글자가 모두 'a'인 문자열의 수이다.
• 문자열에 들어있는 'a'중에서 두 개를 뽑으면, 그 두개를 양 끝으로 하는 부분문자열은 a로 시작하고 a로 끝난다. 그러니까 fa[1]은 문자열에 들어있는 모든 'a'중 에서 두개를 뽑는 경우의 수와 같다. 문자열에 들어있는 모든 'a'의 개수를 ca[1] 이라고 하면 fa[1] = C(ca[1], 2) = ca[1] * (ca[1] - 1) / 2 이다
• 또 문자열에 들어있는 'aa'중에서 두 개를 뽑으면, 그 두개를 양 끝으로 하는 부분문자열은 aa로 시작하고 aa로 끝난다. 그러니까 fa[2]는 비슷하게 fa[2] = C(ca[2], 2) = ca[2] * (ca[2] - 1) / 2 가 된다.
• 이제 식은 간단해졌다. ca값만 다 구해주면, fa를 구할수 있다.
• 전체 문자열을 앞에서부터 스캔하면서 ca를 바로 업데이트하는 것은 비효율적이다. 문자열을 'aaa'까지 읽어서 ca[1]=3, ca[2]=2, ca[3]=1 을 구했는데, 다음 읽은 글자가 또 'a'면, ca[1],ca[2],ca[3],ca[4]를 모두 1씩 증가시켜야 한다. 최악에는 O(n^2)이 걸린다.
• 대신 문자열에 있는 연속된 'a'덩어리가 몇개씩 있는지를 저장한다. 'aaa_aa_aa_a_a_a_a' 라면 count['a']={3:1, 2:2, 1:4} 가 된다. 이것은 O(n)에 가능하다.
• 이제 저 count로부터, ca를 계산할 수 있다. ca[1] = 3*1 + 2*2 + 1*4, ca[2] = (3-1)*1 + (2-1)*2, ca[3] = (3-2)*1 이다.
• 이것은 ca하나를 구하는 데에 O(n)처럼 보일수 있지만, 잘 조작하면 점화식 형태로 바꿀수 있어서, O(1)에 계산 가능하다.
  • t[i] = count[1]+count[2]+..count[i]
  • ca[i+1] = ca[i] - t[i]
  • t[i+1] = t[i] - count[i]
• 이제 모든 ca[0:n_a], 그리고 모든 fa[0:n_a]를 O(n_a)에 구할 수 있다.
• 그리고 모든 글자에 대해서 이걸 계산하는 것도. n_a + n_b + … < n 이므로 결국 O(n) 이다

"""Solution code for "Programmers 68938. 문자열의 아름다움".

- Problem link: https://programmers.co.kr/learn/courses/30/lessons/68938
- Solution link: http://www.teferi.net/ps/problems/programmers/68938
"""

import collections
import itertools


def solution(s):
    # count_by_len_by_ch['x']: counts of 'x', 'xx', 'xxx', ...
    # ex) s: 'ababbab' => count_by_len_by_ch: {['a']: {1:3}, ['b']: {1:2, 2:1}}
    count_by_len_by_ch = collections.defaultdict(
        lambda: collections.defaultdict(int))
    for ch, group in itertools.groupby(s):
        count_by_len_by_ch[ch][len(list(group))] += 1

    # t[i] = counts[i] + counts[i + 1] + ... = t[i - 1] - counts[i - 1]
    # s[i] = i * counts[i] + (i + 1) * counts[i + 1] + ... = s[i - 1] - t[i - 1]
    # f(n) = sum((j - i) for i in range(n) for j in range(i + 1, n))
    #      = (n-1) * n * (n+1) // 6
    # answer = f(len(s)) - sum_ch( sum_l(C(s[l], 2))) )

    answer = (len(s) - 1) * len(s) * (len(s) + 1) // 6
    for count_by_len in count_by_len_by_ch.values():
        t = sum(count_by_len.values())
        s = sum(l * count for l, count in count_by_len.items())
        for l in range(1, max(count_by_len) + 1):
            answer -= s * (s - 1) // 2
            s -= t
            t -= count_by_len[l]
    return answer