#862831 (PyPy3) No.5007 Steiner Space Travel

提出ソース

結果

問題	No.5007 Steiner Space Travel
ユーザー	prussian_coder
提出日時	2023-04-25 14:43:48
言語	PyPy3 (7.3.15)
結果	AC
実行時間	729 ms / 1,000 ms
コード長	13,109 bytes
コンパイル時間	448 ms
コンパイル使用メモリ	87,776 KB
実行使用メモリ	93,252 KB
スコア	8,461,566
最終ジャッジ日時	2023-04-25 14:44:16
合計ジャッジ時間	23,235 ms
ジャッジサーバーID （参考情報）	judge13 / judge12
	純コード判定しない問題か言語

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

テストケース

テストケース表示

入力	結果	実行時間実行使用メモリ
testcase_00	AC	638 ms 91,160 KB
testcase_01	AC	627 ms 89,180 KB
testcase_02	AC	659 ms 91,444 KB
testcase_03	AC	641 ms 90,200 KB
testcase_04	AC	672 ms 90,852 KB
testcase_05	AC	644 ms 90,452 KB
testcase_06	AC	631 ms 90,272 KB
testcase_07	AC	627 ms 89,308 KB
testcase_08	AC	631 ms 91,232 KB
testcase_09	AC	671 ms 91,892 KB
testcase_10	AC	693 ms 91,860 KB
testcase_11	AC	637 ms 91,296 KB
testcase_12	AC	673 ms 92,584 KB
testcase_13	AC	655 ms 91,088 KB
testcase_14	AC	729 ms 93,252 KB
testcase_15	AC	635 ms 90,032 KB
testcase_16	AC	678 ms 91,304 KB
testcase_17	AC	659 ms 91,476 KB
testcase_18	AC	650 ms 90,092 KB
testcase_19	AC	672 ms 91,644 KB
testcase_20	AC	714 ms 91,452 KB
testcase_21	AC	669 ms 90,844 KB
testcase_22	AC	653 ms 91,704 KB
testcase_23	AC	665 ms 90,800 KB
testcase_24	AC	691 ms 91,240 KB
testcase_25	AC	655 ms 89,916 KB
testcase_26	AC	642 ms 91,072 KB
testcase_27	AC	627 ms 91,028 KB
testcase_28	AC	662 ms 90,784 KB
testcase_29	AC	603 ms 89,568 KB

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

ソースコード

raw source code

#from matplotlib import pyplot as plt

def visualize(N,M,pos,center_pos,allocate):
    fig = plt.figure(figsize=(10,10))
    for i in range(N):
        plt.plot(pos[i][0],pos[i][1],color=color_ls[allocate[i]],marker=".")
    for i in range(M):
        plt.plot(center_pos[i][0],center_pos[i][1],color=color_ls[i],marker="*")
    plt.show()
    plt.close()


INF=10**20

import random
from pathlib import Path
import time


LOCAL = False
in_path = "./test"
img_path = "./image"
color_ls = ["red","blue","green","orange","gray","pink","cyan","black"]

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=25): 
    return ((p1[0]-p2[0])**2 + (p1[1]-p2[1])**2) * a

#中心点と各惑星の距離を全探索し、短い順に並びかえる
def calc_dist_from_center(N,M,center_pos,pos):
    dist_list = []
    for m in range(M):
        for n in range(N):
            d = dist(pos[n],center_pos[m])
            dist_list.append((d,m,n))
    return sorted(dist_list)

#中心点から近い惑星順にクラスわけする
def allocate_cluster(N,M,center_pos,pos):
    dist_list = calc_dist_from_center(N,M,center_pos,pos)
    allocate = [-1]*N
    cluster_counts = [0]*M
    total_count = 0
    for d,m,n in dist_list:
        if allocate[n]==-1:
            allocate[n]=m
            cluster_counts[m]+=1
            total_count+=1
        if total_count==N:
            break
    return allocate,cluster_counts
    
#クラスター分けされた点をもとに、k-means法に倣って中心点を算出する。
# グループカウントが多いものは中心点を2つにして、クラスターを分ける。
def calc_center(N,M,pos,allocate,cluster_counts,max_size = 15,split_mode = True):
    point_sums = [[0,0] for _ in range(M)]
    for i in range(N):
        m=allocate[i]
        point_sums[m][0]+=pos[i][0]
        point_sums[m][1]+=pos[i][1]
    center_pos = []

    #clsuter_countが大きい順に見ていき、中心点がM個を超えたら中断する
    if split_mode:
        order = sorted([(cluster_counts[i],i) for i in range(M)],reverse=True)
    else:
        order = [(cluster_counts[i],i) for i in range(M)]
    for _,m in order:  
        if cluster_counts[m]==0:
            center_pos.append([random.randint(0,1000),random.randint(0,1000)])
        else:
            center_pos.append([point_sums[m][0]//cluster_counts[m],point_sums[m][1]//cluster_counts[m]])
            if cluster_counts[m]>=max_size and split_mode:
                x = point_sums[m][0]//cluster_counts[m]
                y = point_sums[m][1]//cluster_counts[m]
                dx = random.randint(-20,20)
                dy = random.randint(-20,20)
                while not (0<=x+dx<=1000 and 0<y+dy<=1000):
                    dx = random.randint(-5,5)
                    dy = random.randint(-5,5)
                center_pos.append([x+dx,y+dy])
        if len(center_pos)==M:
            break

    return center_pos
    
def calc_clustering_distance(center_pos,allocate,pos):
    total_dist = 0
    for i in range(len(pos)):
        m = allocate[i]
        total_dist += dist(pos[i],center_pos[m])
    return total_dist


#K-means法に倣って惑星の点をグループ分けする
def clustering(N,M,pos,max_size):
    center_pos = [[random.randint(0,1000),random.randint(0,1000)] for _ in range(M)] #ランダム初期化
    allocate,cluster_counts= allocate_cluster(N,M,center_pos,pos)
    loop = 0

    while True:
        center_pos = calc_center(N,M,pos,allocate,cluster_counts,max_size=max_size)
        allocate,cluster_counts = allocate_cluster(N,M,center_pos,pos)        
        #print(loop,cluster_counts)
        if loop>=10 and max(cluster_counts)<=max_size:
            break
        loop+=1    

    center_pos = calc_center(N,M,pos,allocate,cluster_counts,split_mode=False)
    cluster_dist = calc_clustering_distance(center_pos,allocate,pos)
    #visualize(N,M,pos,center_pos,allocate)

    return center_pos,allocate,cluster_dist


def f(S,x,n):
    return S*(n+1)+x


#クラスター内にてBitDPでTSPを解く O(n^2*2^n)
def tsp(id_list,pos,center_pos):    
    #point 0~n-1が惑星、nが中心点
    n=len(id_list)-1
    dp = [INF]*(n+1)*(1<<n)
    dp[f(0,n,n)]=0
    for S in range(1<<n):
        for s in range(n+1):
            if dp[f(S,s,n)]==INF:
                continue
            for t in range(n):
                if (S>>t)&1:
                    continue
                S2 = S|(1<<t)
                if s!=n:
                    dp[f(S2,t,n)]=min(dp[f(S,s,n)] + dist(pos[id_list[s]],pos[id_list[t]]),dp[f(S2,t,n)])
                else:
                    dp[f(S2,t,n)]=min(dp[f(S,s,n)] + dist(center_pos,pos[id_list[t]],a=5),dp[f(S2,t,n)])
                dp[f(S2,n,n)]=min(dp[f(S2,t,n)] + dist(center_pos,pos[id_list[t]],a=5),dp[f(S2,n,n)])


    #BitDPから復元            
    path_list = []
    route = []
    state = (1<<n)-1
    now = n
    v = dp[-1]
    e = 10**(-5)
    while state != 0 or now !=n:
        found = False
        if now == n:
            for t in range(n):
                d = dist(center_pos,pos[id_list[t]],a=5)
                if dp[f(state,t,n)]==INF:
                    continue
                if v - dp[f(state,t,n)] >= d - e:
                    if t!=n:
                        route.append(id_list[t])
                    else:
                        path_list.append(route)
                        route = []
                    now = t
                    v -= d
                    found = True
                    break
        else:
            state = state ^ (1<<now)
            for t in range(n+1):
                if dp[f(state,t,n)]==INF:
                    continue
                if t!=n:
                    if not (state>>t)&1:
                        continue
                    d = dist(pos[id_list[now]],pos[id_list[t]])
                else:
                    d = dist(pos[id_list[now]],center_pos,a=5)
                if v - dp[f(state,t,n)] >= d - e:
                    if t!=n:
                        route.append(id_list[t])
                    else:
                        path_list.append(route)
                        route = []
                    now = t
                    v -= d
                    found=True
                    break

    return path_list



def tsp_between_space(center_pos):
    n = 8
    dp = [[INF for _ in range(n)] for _ in range(1<<n)]
    dp[0][-1]=0
    for S in range(1<<n):
        for s in range(n):
            if dp[S][s]==INF:
                continue
            for t in range(n):
                if (S>>t)&1:
                    continue
                S2 = S|(1<<t)
                d = dist(center_pos[s],center_pos[t],a=1)
                dp[S2][t]=min(dp[S][s] + d, dp[S2][t])


    #BitDPから復元            
    state = (1<<n)-1
    s = n-1
    path_list = [s]
    v = dp[-1][-1]
    e = 10**(-5)
    while state != 0:
        state ^=(1<<s)
        for t in range(n):
            if not (state>>t)&1:
                continue
            if dp[state][t]==INF:
                continue
            d = dist(center_pos[s],center_pos[t],a=1)
            if v - dp[state][t] >= d - e:
                path_list.append(t)
                s = t
                v -= d
                break
    return path_list        



def adjust_center_pos(path,pos,M):
    L=len(path)
    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]
        elif path[i]>=0 and path[i+1]<0:
            p,q = -path[i+1]-1,path[i]
        else:
            continue
        edge_count[p]+=1
        edge_x_sum[p]+=pos[q][0]
        edge_y_sum[p]+=pos[q][1]
    center_pos = [[edge_x_sum[m]//edge_count[m],edge_y_sum[m]//edge_count[m]] for m in range(M)]
    return center_pos

def calc_score(path,pos,center_pos):
    score = 0
    L = len(path)
    cost_detail = {1:0,5:0,25:0}
    for i in range(L-1):
        a=1
        if path[i]>=0:
            p1 = pos[path[i]]
            a*=5
        else:
            p1 = center_pos[-path[i]-1]
        if path[i+1]>=0:
            p2 = pos[path[i+1]]
            a*=5
        else:
            p2 = center_pos[-path[i+1]-1]
        score += dist(p1,p2,a)
        cost_detail[a]+=dist(p1,p2,a)
    if LOCAL:
        print(cost_detail)
    return 10**9 / (1000 + score**0.5)

#クラスター内の経路について、前クラスタと最も近い接続点を一番前、後クラスタと最も近い点を一番後ろに持ってくる
def reorder_path(cluster_path,m,pos,prev_pos,next_pos):
    prev_pos_min,next_pos_min = INF,INF
    prev_best,next_best = [],[]
    for l in cluster_path:
        d = dist(prev_pos,pos[l[0]])
        if d<prev_pos_min:
            prev_pos_min=d
            prev_best=l
        d = dist(prev_pos,pos[l[0]])
        if d<prev_pos_min:
            prev_pos_min=d
            prev_best=l[::-1]
    for l in cluster_path:
        if l==prev_best or l==prev_best[::-1]:
            continue
        d = dist(next_pos,pos[l[0]])
        if d<next_pos_min:
            next_pos_min=d
            next_best=l[::-1]
        d = dist(prev_pos,pos[l[0]])
        if d<prev_pos_min:
            next_pos_min=d
            next_best=l
    new_route = prev_best+[-m-1]
    for a in cluster_path:
        if a!=prev_best and a!=prev_best[::-1] and a!=next_best and a!=next_best[::-1]:
            new_route+=a
            new_route+=[-m-1]
    new_route+=next_best

    return new_route


def connect_cluster(p1,p2,m1,m2,m1_id,m2_id):
    d1 = dist(p1,p2,a=25) #そのままつなぐ
    d2 = dist(p1,m1,a=5) + dist(m1,p2,a=5) #m1を経由
    d3 = dist(p1,m2,a=5) + dist(m2,p2,a=5) #m2を経由
    d4 = dist(p1,m1,a=5) + dist(m1,m2,a=1) + dist(m2,p2,a=5) #両方を経由
    d_min = min(d1,d2,d3,d4)
    if d_min == d1:
        return []
    elif d_min == d2:
        return [m1_id]
    elif d_min == d3:
        return [m2_id]
    else:
        return [m1_id,m2_id]



def main(N,M,pos):
    #K-meansで8個に分割
    best_dist = INF
    for _ in range(10):
        center_pos_temp,allocate_temp,cluster_dist = clustering(N,M,pos,50)
        if cluster_dist<best_dist:
            best_dist = cluster_dist
            center_pos = center_pos_temp
            allocate = allocate_temp

    ans = []

    #クラスタ間での回る順番をTSPで決める
    space_order = tsp_between_space(center_pos)

    #クラスタ内で回る順番をTSPで決める
    route = []
    for im,m in enumerate(space_order):
        id_list = [i for i in range(N) if allocate[i]==m]
        n = len(id_list)
        path_in_cluster = []
        #サイズが小さい時は全部使う
        if n<=12:
            path_in_cluster += tsp(id_list+[-m-1],pos,center_pos[m])

        #サイズが大きすぎる時は、clusteringで分割してそれぞれでTSP
        else:
            pos2 = [pos[i] for i in id_list]
            blocks = n//10 + 1
            _,allocate2,_ = clustering(n,blocks,pos2,12)
            for i2 in range(blocks):
                id_list2 = [id_list[i] for i in range(n) if allocate2[i]==i2]
                path_in_cluster += tsp(id_list2+[-m-1],pos,center_pos[m])
        prev_pos = center_pos[space_order[(im-1)%M]]
        next_pos = center_pos[space_order[(im+1)%M]]    
        path_route = reorder_path(path_in_cluster,m,pos,prev_pos,next_pos)
        route.append(path_route)

    #クラスター間の接続方法
    ans = []
    for i in range(M):
        ans += route[i]
        m1_id = space_order[i]
        m2_id = space_order[(i+1)%M]
        ans += connect_cluster(pos[route[i][-1]],pos[route[(i+1)%M][0]],center_pos[m1_id],center_pos[m2_id],-m1_id-1,-m2_id-1)

    #初期位置を0にする
    start_point = ans.index(0)
    ans = ans[start_point:]+ans[:start_point+1] #初期解

    #center-posの位置を補正（最小二乗法の要領）
    center_pos = adjust_center_pos(ans,pos,M)

    score = calc_score(ans,pos,center_pos)
        
    return center_pos,ans,score

def output(center_pos,ans):
    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)
        print(score)
else:
    N,M,pos = read_data("")
    best_score=0
    for _ in range(2):
        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)

yukicoder

結果

テストケース

ソースコード