Atcoder ABC-116 (python)

AtCoder Beginner Contest 116のpythonコード.

• A問題 省略.

• B問題

s = int(input())

a = s
idx = 1
a_list = [a]
while(idx < 1000000):
    idx += 1
    if a % 2 == 0:
        a = int(a / 2)
    else:
        a = 3*a+1
    if a in a_list:
        print(idx)
        break
    else:
        a_list.append(a)

計算時間やばいかと思ったがぎり行けた.

• C問題

import numpy as np
N = int(input())
h_list = list(map(int, input().split()))

def shorten(idx_list, current_h):
    for idx in idx_list:
        current_h[idx] -= 1
    return current_h

def get_highest_idx(current_h):
    max_h = np.max(current_h)
    max_flower_idx = []
    for idx, h in enumerate(current_h):
        if h == max_h:
            max_flower_idx.append(idx)
    return max_flower_idx

ans = 0
while(np.max(h_list) > 0):
    highest_idx_list = get_highest_idx(h_list)
    mizunuki_idx = []
    
    for i, idx in enumerate(highest_idx_list):
        hight = h_list[idx]
        if i==len(highest_idx_list)-1:
            break
        else:
            if idx+1 == highest_idx_list[i+1]:
                j = 0
                while(True):
                    try:
                        if idx+j == highest_idx_list[i+j]:
                            j += 1
                        else:
                            break
                    except:
                        break

                mizunuki = np.array(highest_idx_list)[np.arange(i, i+j,1)]
                mizunuki_idx.append(mizunuki)
    if len(mizunuki_idx) == 0:
        mizunuki_ = [highest_idx_list[0]]
    else:
        mizunuki_ = mizunuki_idx[0]
    
    new_h = shorten(mizunuki_, h_list)
    h_list = new_h
    ans += 1

コードがきたない. 隣り合った一番高い花たちを探して, そこの高さを1減らすのを繰り返す.

• D問題 TLEとREが出ている. 優先度付きキュー(ヒープ)でやるらしいがうまく行かない.

import numpy as np
import heapq
import copy

N, K = list(map(int, input().split()))
t_list, d_list = [], []
for _ in range(N):
    t_, d_ = list(map(int, input().split()))
    t_list.append(t_)
    d_list.append(d_)
    
# Sort by delicoiousness.
c = zip(t_list, d_list)
c = sorted(c, key=lambda x: x[1])[-1::-1]
t_list, d_list = zip(*c)

first_sushi_set = []
nokori_set = []

iter_ = 0
for neta, delicious in zip(t_list, d_list):
    if iter_ < K:
        heapq.heappush(first_sushi_set, (delicious, neta))
    else:
        heapq.heappush(nokori_set, (-delicious, neta))
    iter_ += 1
    
first_neta_num = len(list(set(t_list[:K])))

def get_next_sushi_set(current_sushi, nokori_sushi):
    current_sushi_cp = copy.deepcopy(current_sushi)

    while(len(current_sushi) > 0):
        #種類数が減らないかつ一番美味しくない寿司を選ぶ
        remove_sushi = heapq.heappop(current_sushi)
        score, neta = remove_sushi
        if len(current_sushi) == 0:
            break
        if neta in np.array(current_sushi)[:, 1]:
            break
        else:
            pass
        
    current_sushi_cp.remove(remove_sushi)
    heapq.heappush(nokori_sushi, (-score, neta))
    nokori_sushi_cp = copy.deepcopy(nokori_sushi)

    while(len(nokori_sushi) > 0):
        #種類数が増えるかつ一番美味しい寿司を選ぶ
        add_sushi = heapq.heappop(nokori_sushi)
        negative_score, neta = add_sushi
        score = - negative_score
        if len(nokori_sushi) == 0:
            break
        if neta in np.array(current_sushi_cp)[:, 1]:
            pass
        else:
            break        
    heapq.heappush(current_sushi_cp, (score, neta))
    nokori_sushi_cp.remove(add_sushi)
            
    return current_sushi_cp, nokori_sushi_cp
        
current_sushi_set = first_sushi_set
nokori_sushi_set = nokori_set
score_list = []
for neta_num in range(first_neta_num, N+1, 1):
    score =  len(set(np.array(current_sushi_set)[:, 1]))**2 + sum(np.array(current_sushi_set)[:, 0])
    score_list.append(score)
    current_sushi_set, nokori_sushi_set = get_next_sushi_set(current_sushi_set, nokori_sushi_set)
print(max(score_list))