For Small Dataset

The following should work fine for small data set, even if without the dynamic programming.

import statistics

def mean(numbers):
    if len(numbers) == 0:
        return 0
    return statistics.mean(numbers)

def dip(numbers, time, d): 
    if time == 0:
        d[0] = mean(numbers)
        return d[0]
    
    exp_val = d[time-1] if (time - 1) in d else dip(numbers, time - 1, d)
    larger = [i for i in numbers if i > exp_val]
    p = len(larger) / len(numbers)
    return mean(larger) * p +  exp_val * (1 - p)


N = int(input())
for n in range(N):
    [count, time] = [int(i) for i in input().split()]
    numbers = [int(i) for i in input().split()]
    d = dict()
    res = dip(numbers, time, d)
    print("Case #{}: {:.6f}".format(n + 1, res))

For Large Dataset

The most consuming part in the previous one is to find out the “larger” numbers, we find them all in the small dataset solution. But when the dataset grows larger, that does not work. To make that easier, we can sort the data once, and binary sort afterwards. So that the computation can be reduced to O(Nlg(N) + N + KlgN) = O((N + K)lgN).

Failed Attempt

What I got was always an RE error, maybe I should try later.

import statistics
from math import floor

def mean(numbers):
    if len(numbers) == 0:
        return 0
    return statistics.mean(numbers)

def larger_index_utility(numbers, number, front, end):
    if front >= end:
        return front

    mid = floor((front + end) / 2)

    if number < numbers[mid]:
        return larger_index_utility(numbers, number, front, mid - 1)
    if number == numbers[mid]:
        return mid
    return larger_index_utility(numbers, number, mid + 1, end)

def larger_index(numbers, number):
    index = larger_index_utility(numbers, number, 0, len(numbers) - 1)
    if numbers[index] >= number:
        return index
    return index + 1


def dip(numbers, time, d, suffix): 
    if time == 0:
        d[0] = mean(numbers)
        return d[0]
    
    if (time - 1) in d:
        exp_val = d[time - 1]
    else:
        exp_val =  dip(numbers, time - 1, d, suffix)
        d[time - 1] = exp_val
    index = larger_index(numbers, exp_val)
    larger = numbers[index:]
    p = 1 - len(larger) / len(numbers)
    if index in suffix:
        suffix_sum = suffix[index]
    else:
        suffix_sum = sum(numbers[index:])
        suffix[index] = suffix_sum
    return suffix_sum / len(numbers) +  exp_val * p


N = int(input())
for n in range(N):
    [count, time] = [int(i) for i in input().split()]
    numbers = [int(i) for i in input().split()]
    d = dict()
    suffix = dict()
    numbers.sort()
    res = dip(numbers, time, d, suffix)
    print("Case #{}: {:.6f}".format(n + 1, res))

Successful Attempt

I found that the range is specified as 0 ≤ K ≤ 5 * 10⁴ , so the number of recursion frames can be as large as this one. If I write this in a recursive way, it is bounded to overflow. So I finally wrote it in an iterative manner.

#include <cstdio>
#include <algorithm>
#include <vector>
using namespace std;

int a[20010];
long long sum[20010];
double dp[50010];

int main()
{
    int T, N, K;
    scanf("%d", &T);

    int iCase = 0;
    while(T--) 
    {
        iCase++;
        scanf("%d %d", &N, &K);
        for(int i = 0; i < N; i++)
        {
            scanf("%d", &a[i]);
        }
        sort(a, a+N);
        sum[N] = 0;
        for(int i = N - 1; i >= 0; i--)
        {
            sum[i] = sum[i+1] + a[i];
        }
        dp[0] = (double)sum[0] / N;
        for(int i = 1; i <= K; i++)
        {
            int cnt = upper_bound(a, a+N, dp[i-1]) - a;
            dp[i] = (double)(dp[i - 1] * cnt + sum[cnt]) / N;
        }
        printf("Case #%d: %f\n", iCase, dp[K]);

    }
    return 0;
}

C++ STL Review

• std::sort()
• std::upper_bound()

These 2 are STL functions, but also can be applied to C-style arrays. For the upper_bound function, the type of the return value is a pointer, to extract the index in integer, all need to be done is to extract the starting element from it., as in the following,

 int upper1 = upper_bound(arr, arr+5, 35) - arr;