[ココア出力]ランキング検索/Python/資料構造と探索

ランキング検索

質問する

せいげんじょうけん

解答方法

from itertools import combinations as combi
from collections import defaultdict

def solution(infos, queries):
    answer = []
    info_dict = defaultdict(list)
    for info in infos:
        info = info.split()
        info_key = info[:-1]
        info_val = int(info[-1])
        for i in range(5): # 4가지의 정보에 대해 16가지의 조합
            for c in combi(info_key, i):
                tmp_key = "".join(c)
                info_dict[tmp_key].append(info_val)
                
    for key in info_dict.keys(): # 각 value(코딩테스트 점수)를 오름차순 정렬
        info_dict[key].sort()
        
    for query in queries:
        query = query.split(' ')
        query_score = int(query[-1])
        query = query[:-1]
        
        tmp_q = ''.join(query)
        tmp_q = tmp_q.replace('and', '')
        tmp_q = tmp_q.replace('-', '')
        
        # lower bound(원하는 값 이상이 처음 나오는 위치를 탐색)
        if tmp_q in info_dict:
            scores = info_dict[tmp_q]
            if len(scores) > 0:
                start, end = 0, len(scores)
                # binary search
                while end > start:
                    mid = (start + end) // 2
                    if scores[mid] >= query_score:
                        end = mid
                    else:
                        start = mid + 1
                answer.append(len(scores) - start)
        else:
            answer.append(0)
        
                
    return answer