목차

문자열의 아름다움

ps
링크programmers.co.kr/…
출처프로그래머스
문제 번호68938
문제명문자열의 아름다움
레벨Level 5
분류

애드혹

시간복잡도O(n)
인풋사이즈n<=300,000
사용한 언어Python
해결날짜2021/01/25

풀이

코드

"""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