Прошу помочь с решением задачи со старого Яндекс Контеста

В качестве тренировки решал тренировочный контест от Яндекса. Смог решить все задачи, кроме одной. К сожалению, ни поиск в интернете, ни помощь нейросети не помогли.

Итак, вот условия самой задачи:

Вася взял игральную кость и написал на гранях числа a1, a2, a3, a4, a5, a6.

Для генерации случайного числа Вася решил воспользоваться следующим алгоритмом:

  • Выбрать число k.
  • Подбросить кубик k раз и записать на листик последовательно выпавших чисел bj.
  • Пройтись по списку с конца и вычеркнуть число bj, если оно равно b[j−1] (b1 всегда останется в последовательности).

Определите математическое ожидание суммы оставшихся в последовательности чисел, если Вася сообщит вам числа ai, k. (1≤ai≤1000), (1≤k≤1000).

Обратите внимание, что кубик у Васи честный и все выпадение любой из граней равновероятно. Кроме этого, подбрасывания кубика независимы.​

Ответ будет считаться верным, если относительная или абсолютная погрешность не будет превышать 10^−6

Пример 1

Ввод
1 2 3 4 5 6 2

Вывод 6.4166666667

Пример 2

Ввод
1 1 1 1 1 1 3

Вывод 1.0000000000

Пример 3

Ввод
1 2 1 2 2 2 2

Вывод 2.3333333333

Подход к решению

Считать в лоб все возможные суммы и их вероятности или прямо набирать статистику, имитируя Васин алгоритм, при данных ограничениях совершенно бесполезно. На данный момент, единственное, что приходит в голову это вычислить E(S) как E(N)*E(a), где E(N) мат ожидание количества оставшихся чисел, E(a) мат ожидание bj.

Такой подход работает при классическом кубике 1, 2, 3, 4, 5, 6, но как его обобщить на произвольные ai я не понимаю. Возможно, я что-то упускаю или есть какое-то более програмисткое решение, которое работает.

Вот моя реализация описанного подхода:

from collections import Counter

with open('input.txt', 'r') as input:
    nums = list(map(int, input.readline().split()))
    k = int(input.readline())

counter = Counter(nums)
prob_nums = [v / 6 for v in counter.values()] # p(ai)
E_a = sum(nums) / 6 # мат ожидание значения bj
prob_dupl = sum(map(lambda x: x**2, prob_nums)) # p(bj == b[j-1])
E_n = 1 + (k - 1) * (1 - prob_dupl) # мат ожидание количества оставшихся bj

print(E_a * E_n)

Ответы (1 шт):

Автор решения: Stanislav Volodarskiy

У задачи есть красивое решение. Но я умею добраться до него только с третьей попытки.

1

Обозначим sk,a сумму всех чисел во всех (уже сокращенных) комбинациях из k чисел, которые оканчиваются на число a. Количество таких комбинаций равно 6k-1 - последнее число фиксированно, каждое из k-1 предыдущих выбирается из шести кандидатов.

Ответ задачи можно выразить так: [∑a sk,a] / 6k. Сумма ведётся по всем значениям a, то есть мы перебираем все суммы всех комбинаций. Если разделить сумму на общее количество, получится нужное среднее.

NB Если число на кубике повторяется несколько раз, оно участвует в сумме несколько раз.

sk+1,a = ∑b≠a [sk,b + 6k-1a] + ∑b=a sk,b

Сумма b≠a ... ведётся по всем граням кубика, значение которых не равно a. Сумма b=a ... ведётся по всем граням кубика, значение которых равно a.

Первое слагаемое складывает те комбинации, которые увеличились в длину с добавлением a. Второе слагаемое складывает комбинации, которые уже завершались на a и, следовательно, их сумма не изменилась.

Сгруппируем суммы по другому:

sk+1,a = ∑b≠a 6k-1a + ∑b sk,b

Рекурентная формула есть, добавим начальное условие: s1,a = a. Сумма комбинаций из одного числа сопадает с этим самым числом.

Пора писать код:

*a, k = map(int, input().split())

c = 1
s = a
for _ in range(1, k):
    s_next = [0] * len(a)
    for i in range(len(a)):
        for j in range(len(a)):
            s_next[i] += s[j]
            if a[j] != a[i]:
                s_next[i] += a[i] * c
    c *= len(a)
    s = s_next

print(sum(s) / (len(a) * c))

2

Введём tk,a = sk,a/6k. Тогда

ответ задачи есть a tk,a,

реккурентная формула tk+1,a = ∑b≠a a/36 + ∑b tk,b/6,

начальное условие t1,a = a/6.

Пишем код:

*a, k = map(int, input().split())

t = [ai / len(a) for ai in a]
for _ in range(1, k):
    t_next = [0] * len(a)
    for i in range(len(a)):
        for j in range(len(a)):
            t_next[i] += t[j] / len(a)
            if a[j] != a[i]:
                t_next[i] += a[i] / (len(a) * len(a))
    t = t_next

print(sum(t))

3

Наденем на реккурентную формулу сумму a ...: a tk+1,a = ∑a [∑b≠a a/36 + ∑b tk,b/6].

Второе слагаемое от a не зависит. Получаем: a tk+1,a = ∑ab≠a a/36 + ∑b tk,b

Обозначим Tk = ∑a tk,a. Тогда Tk+1 = ∑ab≠a a/36 + Tk.

Первое слагаемое - константа, а Tk - арифметическая прогрессия:

начальное значение T1 = ∑a a/6,

шаг прогрессии ΔT = ∑ab≠a a/36,

общая формула Tk = T1 + (k - 1)ΔT.

Пишем код:

*a, k = map(int, input().split())

dt = 0
for ai in a:
    for aj in a:
        if aj != ai:
            dt += ai
dt /= len(a) * len(a)

t1 = sum(a) / len(a)

print(t1 + (k - 1) * dt)

Задача решена за время, которое не зависит от k.

→ Ссылка