Python Dictionary/Hashmap Search Optimization Problem-CodePudding

I am trying to solve this problem with Python, and it has a CPU time limit of 1 second and a memory limit of 1024 megabytes: https://open.kattis.com/contests/v4xrif/problems/grandpabernie

Over the years, Grandpa Bernie has traveled all over the world. He doesn’t travel that much anymore, but he loves to tell his grandchildren stories from all these trips. He’ll tell them the story from when he went to Israel for the first time, or when he went to Greece for the third time.

His memory works in a funny way. He can easily remember his k:th trip to a particular country, but he’ll have a hard time remembering in which year he went on that trip. Given a list of all the trips Grandpa Bernie went on, can you answer a number of queries asking in which year he went on his k:th trip to a particular country?

Input The input consists of:

one line with one integer n (1≤n≤100000), the number of trips Grandpa Bernie went on;

n lines each containing the name s (1≤|s|≤20) of a country and an integer y (1≤y≤1000000) representing a trip to country s that Grandpa Bernie went on in year y;

one line with one integer q (1≤q≤100000), the number of queries;

q lines each containing the name s of a country and an integer k representing a query for the k:th time Grandpa Bernie went to country s.

Each country name only consists of letters from the English alphabet. It is also guaranteed that, for each query asking for the k:th trip to country s, k is at least 1 and no greater than the number of times Grandpa Bernie went to country s. In particular, it is guaranteed that Grandpa Bernie has visited country s at least once.

Output For each query for the k:th trip Grandpa Bernie went to a country s, output a single line containing the year in which Grandpa Bernie went on that trip.

Sample Input 1

4
Iceland 2016
Sweden 2015
Iceland 1982
Norway 1999
3
Sweden 1
Iceland 1
Iceland 2

Sample Output 1

2015
1982
2016

Sample Input 2

4
Iceland 2014
Iceland 2015
Iceland 2015
Iceland 2016
4
Iceland 4
Iceland 3
Iceland 2
Iceland 1

Sample Output 2

My strategy for this problem would be to create a dictionary/hashmap with the country name as the keys, and the dates he has gone to the country as a list of values. Then, I would go through each query and get the value from the dictionary/hashmap that corresponds to the query.

My code:

dates = {}
for i in range(int(input().strip())):
    line = input().strip().split()
    city = line[0]
    year = int(line[1])
    if city in dates:
        dates[city].append(year)
    else:
        dates[city] = [year]

for i in range(int(input().strip())):
    line = input().strip().split()
    print(sorted(dates[line[0]])[int(line[1])-1])

As you can see, this is a relatively short program, but the speed is not fast enough. Is there a way to speed it up?

CodePudding user response：

The slowest part of you program is going to be the repeated sorting each time you query.

To speed things up a bit-

from collections import defaultdict

trips = defaultdict(list)

num_trips = int(input())
for _ in range(num_trips):
    destination, year = input().strip().split()
    trips[destination].append(year)

for years in trips.values():
    years.sort()

num_queries = int(input())
for _ in range(num_queries):
    destination, occurance = input().strip().split()
    year = trips[destination][int(occurance)-1]
    print(year)

Additionally, you could delay sorting the years until a query arrives for that country, and keep a list of already-sorted countries, however unless you have an input that gives you a LOT of entries in one country and then never queries that country, you won't see a speedup.

If you know that the query list isn't too long, you can also preprocess it to figure out in advance a unique set of countries that need sorting, but that very well could not be the case, causing you to exceed the memory limit.