google kickstart round f. #2 Festival

https://codingcompetitions.withgoogle.com/kickstart/round/0000000000435bae/0000000000887dba

Problem

You have just heard about a wonderful festival that will last for D days, numbered from 1 to D. There will be N attractions at the festival. The i-th attraction has a happiness rating of hi and will be available from day si until day ei, inclusive.

You plan to choose one of the days to attend the festival. On that day, you will choose up to K attractions to ride. Your total happiness will be the sum of happiness ratings of the attractions you chose to ride.

What is the maximum total happiness you could achieve?

Input

The first line of the input gives the number of test cases, T. T test cases follow.

The first line of each test case contains the three integers, D, N and K. The next N lines describe the attractions. The i-th line contains hi, si and ei.

Output

For each test case, output one line containing Case #x: y, where x is the test case number (starting from 1) and y is the maximum total happiness you could achieve.

Limits

Memory limit: 1 GB.
1≤T≤100.
1≤K≤N.
1≤si≤ei≤D, for all i.
1≤hi≤3×105, for all i.
Test Set 1
Time limit: 20 seconds.
1≤N≤1000.
1≤D≤1000.
Test Set 2
Time limit: 90 seconds.
For at most 10 test cases:

1≤N≤3×105.
1≤D≤3×105.

For the remaining cases, 1≤N,D≤1000.

Sample

2
10 4 2 
800 2 8
1500 6 9
200 4 7
400 3 5
5 3 3
400 1 3
500 5 5
300 2 3

Case #1: 2300
Case #2: 700

(Naive) Solution

Below code passes for Test set 1 but exceeds time limit for Test set 2.

def f(D, N, K, attractions):
    """Returns the maximum total happiness
    
    When given the attractions up to `K` in the festival,
    it returns the maximum total happiness with one day attention.
    
    Args:
        D (int): festival days
        N (int): the number of attractions at the festival
        K (int): the number of attractions can be chosen
        attractions ([[hi, si, ei]]): 
            hi: happiness rating,  
            si: available days from `si` 
            ei: available days up to `ei` (inclusive)
    
    Returns:
        happiness(int): the maximum total happiness chosen at the festival
    """
    def make_attraction_map(attractions):
        ret = {}
        for attraction in attractions:
            hi, si, ei = attraction
            for day in range(si, ei+1):
                ret[day] = [hi] + ret.get(day, [])
        return ret
                
    max_happiness = 0
    # key: day, value: list of happiness of possible attractions
    attraction_map = make_attraction_map(attractions)
    
    for i in range(1, D+1):
        if not i in attraction_map:
            continue
        attractions = attraction_map[i]
        happiness = sum(sorted(attractions, reverse=True)[:K])
        max_happiness = max(max_happiness, happiness) 
    
    return max_happiness

def read_ints(): return [int(v) for v in input().strip().split()]

if __name__ == '__main__':
    T = int(input())
    for i in range(T):
        D, N, K = read_ints()
        # [[hi, si, ei]] where `hi` happiness rating, `si` <= available days <= `ei`
        attractions = [read_ints() for i in range(N)]
        print (f'Case #{i+1}: {f(D, N, K, attractions)}')

Optimized Solution

TBD