#863431 (PyPy3) No.5007 Steiner Space Travel

提出ソース

結果

問題	No.5007 Steiner Space Travel
ユーザー	prussian_coder
提出日時	2023-04-28 11:52:02
言語	PyPy3 (7.3.15)
結果	AC
実行時間	964 ms / 1,000 ms
コード長	13,299 bytes
コンパイル時間	639 ms
コンパイル使用メモリ	87,356 KB
実行使用メモリ	94,476 KB
スコア	8,549,161
最終ジャッジ日時	2023-04-28 11:52:36
合計ジャッジ時間	31,995 ms
ジャッジサーバーID （参考情報）	judge15 / judge13
	純コード判定しない問題か言語

このコードへのチャレンジ
（要ログイン）

テストケース

テストケース表示

入力	結果	実行時間実行使用メモリ
testcase_00	AC	946 ms 93,372 KB
testcase_01	AC	949 ms 94,020 KB
testcase_02	AC	959 ms 92,664 KB
testcase_03	AC	948 ms 93,784 KB
testcase_04	AC	951 ms 92,128 KB
testcase_05	AC	942 ms 93,432 KB
testcase_06	AC	947 ms 92,264 KB
testcase_07	AC	961 ms 94,216 KB
testcase_08	AC	947 ms 93,164 KB
testcase_09	AC	948 ms 92,416 KB
testcase_10	AC	948 ms 93,280 KB
testcase_11	AC	953 ms 92,844 KB
testcase_12	AC	948 ms 93,348 KB
testcase_13	AC	944 ms 92,652 KB
testcase_14	AC	949 ms 94,476 KB
testcase_15	AC	957 ms 92,848 KB
testcase_16	AC	958 ms 93,028 KB
testcase_17	AC	949 ms 92,568 KB
testcase_18	AC	940 ms 91,812 KB
testcase_19	AC	944 ms 91,520 KB
testcase_20	AC	957 ms 92,784 KB
testcase_21	AC	964 ms 91,108 KB
testcase_22	AC	947 ms 92,256 KB
testcase_23	AC	950 ms 93,328 KB
testcase_24	AC	944 ms 92,888 KB
testcase_25	AC	954 ms 92,620 KB
testcase_26	AC	944 ms 93,600 KB
testcase_27	AC	943 ms 92,708 KB
testcase_28	AC	945 ms 92,160 KB
testcase_29	AC	957 ms 92,028 KB

権限があれば一括ダウンロードができます

ソースコード

raw source code

import random
from pathlib import Path
import time
import os
import math

LOCAL = False
in_path = "./test"
out_path = "./test/result"
INF=10**20
alpha = 5
FILE_OUTPUT = True

def read_data(file):
    if LOCAL:
        with open(file,mode="r") as f:
            data = f.readlines()
        N,M = map(int,data[0].split())
        pos = [[int(x) for x in data[i+1].split()] for i in range(N)]
    else:
        N,M=map(int,input().split())
        pos = [[int(x) for x in input().split()] for i in range(N)]
    return N,M,pos

#2点間の距離を返す 
def dist(p1,p2,a=alpha**2): 
    return ((p1[0]-p2[0])**2 + (p1[1]-p2[1])**2) * a

# 一番近い惑星を貪欲に選び続ける（Nearest Neighbour法）
def nearest_neighbor(pos,N):
    start_point = random.randint(0,N-1) #初期位置はランダム
    v = start_point
    visited = [False] * N
    visited[v] = True
    route = [v]

    # 初期惑星以外のN-1個の惑星を訪問していく
    for _ in range(N - 1):
        nearest_dist = INF
        nearest_v = -1

        # 一番近い惑星を探す
        for next in range(N):
            if visited[next]:
                continue

            d = dist(pos[v], pos[next])
            if d < nearest_dist:
                nearest_dist = d
                nearest_v = next

        # 次の頂点に移動
        v = nearest_v
        visited[v] = True
        route.append(nearest_v)

    return route

#焼きなまし関数
def simulated_annealing(score,temp):
    if score<=0:
        return True
    elif score/temp > 10:
        return False
    return math.exp(-score/temp) > random.random()

#2-opt近傍
def two_opt(state,temp):
    #print("two-opt")
    i = random.randint(1,state.N-2)
    j = random.randint(1,state.N-2)
    if i>j:
        i,j=j,i
    #(a,b)-(c,d)->(a,c)-(b-d)
    a = state.route[i]
    b = state.route[i+1]
    c = state.route[j]
    d = state.route[j+1]
    if len(set({a,b,c,d}))!=4:
        return False
    ab_cd = state.calc_distance(a,b) + state.calc_distance(c,d)
    ac_bd = state.calc_distance(a,c) + state.calc_distance(b,d)
    if simulated_annealing(ac_bd-ab_cd,temp):     
        state.route[i+1:j+1] = state.route[i+1:j+1][::-1]
        state.cost += ac_bd-ab_cd
        return True
    else:
        return False

#挿入近傍
def point_insert(state,temp):
    #print("Point_insert")
    i = random.randint(1,state.N-2)
    j = random.randint(1,state.N-2)
    if i>j:
        i,j=j,i

    #(a,b,c)-(d,e)->(a,c)-(d,b,e)
    a = state.route[i-1]
    b = state.route[i]
    c = state.route[i+1]
    d = state.route[j]
    e = state.route[j+1]
    if len(set({a,b,c,d,e}))!=5:
        return False
    abc_de = state.calc_distance(a,b) + state.calc_distance(b,c) + state.calc_distance(d,e)
    ac_dbe = state.calc_distance(a,c) + state.calc_distance(d,b) + state.calc_distance(b,e)
    if simulated_annealing(ac_dbe-abc_de,temp):
        state.route.remove(b)
        if i<j:
            state.route.insert(j,b)
        else:
            state.route.insert(j+1,b)
        state.cost += ac_dbe-abc_de
        return True
    else:
        return False

#3-opt
def three_opt(state,temp):
    #print("three-opt")

    i = random.randint(1,state.N-2)
    j = random.randint(1,state.N-2)
    k = random.randint(1,state.N-2)
    i,j,k = sorted((i,j,k))
    if abs(i-j)<=1 or abs(j-k)<=1:
        return False

    cost_list = []
    a, b, c, d, e, f = state.route[i - 1], state.route[i], state.route[j - 1], state.route[j], state.route[k - 1], state.route[k % len(state.route)]
    current_cost = state.calc_distance(a,b) + state.calc_distance(c,d) + state.calc_distance(e,f)
    for mode in range(8):
        A,B,C,D,E,F = a,b,c,d,e,f
        if (mode>>2)&1:
            B,C=C,B
        if (mode>>1)&1:
            D,E=E,D
        if mode&1:
            F,A=A,F
        new_cost = state.calc_distance(A,B) + state.calc_distance(C,D) + state.calc_distance(E,F)
        cost_list.append((new_cost-current_cost,mode))
    cost_list.sort()
    if cost_list[0][1]!=0:
        next_mode = cost_list[0][1]
        state.cost += cost_list[0][0]
    elif simulated_annealing(cost_list[1][0],temp):
        next_mode = cost_list[1][1]
        state.cost += cost_list[1][0]
    else:
        return False
    
    if (i - 1) < (k % len(state.route)):
        first_segment = state.route[k% len(state.route):] + state.route[:i]
    else:
        first_segment = state.route[k % len(state.route):i]
    second_segment = state.route[i:j]
    third_segment = state.route[j:k]    
    if next_mode&1:
        first_segment.reverse()
    if (next_mode>>1)&1:
        third_segment.reverse()
    if (next_mode>>2)&1:
        second_segment.reverse()
    state.route = first_segment + second_segment + third_segment
    return True

#宇宙船の位置を少し動かす
def move_station(state,temp):
    move_range = 20
    m = random.randint(0,state.M-1)
    x,y = state.center_pos[m]
    dx = random.randint(-move_range,move_range)
    dy = random.randint(-move_range,move_range)
    current_cost = state.cost
    x_new = max(0,min(1000,x+dx))
    y_new = max(0,min(1000,y+dy))
    state.center_pos[m] = [x_new,y_new]
    new_cost = state.calc_cost()
    if simulated_annealing(new_cost-current_cost,temp):
        return True
    else:
        state.center_pos[m]=[x,y]

def adjust_station(state,temp):
    path = state.ans()
    L=len(path)
    M=state.M
    edge_count = [0]*M
    edge_x_sum = [0]*M
    edge_y_sum = [0]*M
    for i in range(L-1):        
        if path[i]<0 and path[i+1]>=0:
            p,q = -path[i]-1,path[i+1]
            edge_count[p]+=5
            edge_x_sum[p]+=pos[q][0]*5
            edge_y_sum[p]+=pos[q][1]*5

        elif path[i]>=0 and path[i+1]<0:
            p,q = -path[i+1]-1,path[i]
            edge_count[p]+=5
            edge_x_sum[p]+=pos[q][0]*5
            edge_y_sum[p]+=pos[q][1]*5

        elif path[i]<0 and path[i+1]<0:
            p,q = -path[i]-1,-path[i+1]-1
            edge_count[p]+=1
            edge_count[q]+=1
            edge_x_sum[p]+=pos[q][0]
            edge_y_sum[p]+=pos[q][1]
            edge_x_sum[q]+=pos[p][0]
            edge_y_sum[q]+=pos[p][1]
    for m in range(M):
        if edge_count[m]==0:
            state.center_pos[m] = [random.randint(0,1000),random.randint(0,1000)]
        else:
            state.center_pos[m] = [edge_x_sum[m]//edge_count[m],edge_y_sum[m]//edge_count[m]]
    state.cost = state.calc_cost()
    return True        


#焼きなましで経路を短くする            
def optimize_route(state,mode_ls,time_limit,start_temp,end_temp):
    dt=time.time()
    trial = dict()
    success = dict()
    for mode in mode_ls:
        trial[mode[0]]=0
        success[mode[0]]=0
    while time.time()-dt<time_limit:
        temp = start_temp + (end_temp - start_temp) * (time.time()-dt)/time_limit
        rand = random.random()
        for mode in mode_ls:
            if rand<mode[1]:
                
                trial[mode[0]]+=1
                flag = mode[0](state,temp)
                if flag:
                    success[mode[0]]+=1
                break
    #print(trial,success)



#隣接する惑星間の距離が長い惑星を出力する O(N)
# limit: 出力する個数        
def find_edge_point(state,limit):
    edge_list = [(state.calc_distance(state.route[i-1],state.route[i]),state.route[i-1],state.route[i]) for i in range(len(state.route))]
    edge_list.sort(reverse=True)
    point_ls = set()
    count = 0
    for v,i,j in edge_list:
        if not i in point_ls:
            point_ls.add(i)
            count+=1
        if not j in point_ls:
            point_ls.add(j)
            count+=1
        if count >= limit:
            return point_ls
        
def K_means_clustering(state):
    use_list = find_edge_point(state,30)
    valid_station = {0,1,2,3,4}
    for m in range(state.M):
        if not m in valid_station:
            state.center_pos[m]=[INF,INF]
    for _ in range(40):
        state.find_nearest_station()
        state.set_station(use_list,valid_station)
        
    use_list = find_edge_point(state,30)
    valid_station = {0,1,2,3,4,5,6,7}
    for m in range(state.M):
        if not m in valid_station:
            state.center_pos[m]=[INF,INF]
    for _ in range(40):
        state.find_nearest_station()
        state.set_station(use_list,valid_station)


class State:
    def __init__(self,N,M,pos,initial_route):
        self.N=N
        self.M=M
        self.route = initial_route(pos,N)
        self.allocate = [-1]*N
        self.pos = pos
        self.center_pos = [[random.randint(0,1000),random.randint(0,1000)] for _ in range(M)]
        self.cost = self.calc_cost()

    #惑星iと惑星jの移動について、宇宙ステーションを用いた際の最適な距離とルートを返す
    def calc_distance(self,i,j,output_path = False):
        path_candidate = [] #(distance,connector)       
        #p1→p2
        d = dist(self.pos[i],self.pos[j])
        path_candidate.append((d,[]))

        #p1→m1→p2
        if self.allocate[i]!=-1:
            m1 = self.allocate[i]
            d = dist(self.pos[i],self.center_pos[m1],a=alpha) + dist(self.pos[j],self.center_pos[m1],a=alpha)
            path_candidate.append((d,[-m1-1]))

        #p1→m2→p2
        if self.allocate[j]!=-1:
            m2 = self.allocate[j]
            d = dist(self.pos[i],self.center_pos[m2],a=alpha) + dist(self.pos[j],self.center_pos[m2],a=alpha)
            path_candidate.append((d,[-m2-1]))

        #p1→m1→→m2→p2
        if self.allocate[i]!=-1 and self.allocate[j]!=-1:
            m1 = self.allocate[i]
            m2 = self.allocate[j]
            d = dist(self.pos[i],self.center_pos[m1],a=alpha) + dist(self.pos[j],self.center_pos[m2],a=alpha) + dist(self.center_pos[m1],self.center_pos[m2],a=1)
            path_candidate.append((d,[-m1-1,-m2-1]))
        path_candidate.sort()
        if output_path:
            return path_candidate[0][1]
        else:
            return path_candidate[0][0]
    
    #1周のコストを計算する
    def calc_cost(self):
        cost = 0
        for i in range(len(self.route)):
            cost += self.calc_distance(self.route[i-1],self.route[i])      
        return cost
    
    def ans(self):
        res = []
        start_point = 0
        p1 = start_point
        v1 = self.route.index(start_point)
        for _ in range(self.N):
            res.append(p1)
            v2 = (v1+1)%self.N
            p2 = self.route[v2]
            connector = self.calc_distance(p1,p2,output_path=True)
            for c in connector:
                res.append(c)
            v1 = v2
            p1 = p2
        res.append(p1)
        return res
    
    def find_nearest_station(self):
        for i in range(self.N):
            d = [dist(self.pos[i],self.center_pos[j]) for j in range(self.M)]
            self.allocate[i] = d.index(min(d))

    def set_station(self,use_list,valid_station):
        point_count = [0]*M
        point_pos_sum = [[0,0] for _ in range(M)]
        for v in use_list:
            m = self.allocate[v]
            point_count[m]+=1
            point_pos_sum[m][0]+=self.pos[v][0]
            point_pos_sum[m][1]+=self.pos[v][1]
        for m in range(M):
            if not m in valid_station:
                continue
            if point_count[m]==0:
                self.center_pos[m]=[random.randint(0,1000),random.randint(0,1000)]
            else:
                self.center_pos[m]=[point_pos_sum[m][0]//point_count[m],point_pos_sum[m][1]//point_count[m]]
    


        





def main(N,M,pos):
    state = State(N,M,pos,nearest_neighbor)    
    mode_ls = [(two_opt,0.6),(three_opt,0.8),(point_insert,1)]
    optimize_route(state,mode_ls,0.15,500,100)
    K_means_clustering(state)
    mode_ls = [(adjust_station,0.005),(move_station,0.3),(two_opt,0.7),(three_opt,0.85),(point_insert,1)]
    optimize_route(state,mode_ls,0.52,500,100)
    adjust_station(state,0)
    return state.center_pos,state.ans(),state.cost

def output(center_pos,ans,score,file=""):
    if LOCAL and FILE_OUTPUT:
        Path(out_path).mkdir(exist_ok=True)
        file_out = os.path.join(out_path,file.stem + "_" + str(int(score))+".txt")
        with open(file_out,mode="w") as f:
            for x,y in center_pos:
                f.write(str(x)+" "+str(y)+"\n")
            f.write(str(len(ans))+"\n")
            for a in ans:
                if a<0:
                    f.write("{0} {1}\n".format(2,-a))
                else:
                    f.write("{0} {1}\n".format(1,a+1))
    else:
        for x,y in center_pos:
            print(x,y)
        print(len(ans))
        for a in ans:
            if a<0:
                print(2,-a)
            else:
                print(1,a+1)


if LOCAL:
    file_ls = Path(in_path).glob("*.txt")
    for file in file_ls:
        print(file)
        N,M,pos = read_data(file)
        center_pos,ans,score = main(N,M,pos)
        output(center_pos,ans,score,file)
        print(score)
else:
    N,M,pos = read_data("")
    best_score=0
    for _ in range(1):
        center_pos,ans,score = main(N,M,pos)
        if score>best_score:
            best_score = score
            best_pos = center_pos
            best_ans = ans
    output(best_pos,best_ans,score)

yukicoder

結果

テストケース

ソースコード