Submission #21632160


Source Code Expand

"""
import sys
import copy
import time
import numpy as np
import numba
from numba import njit, b1, i1, i4, i8, f8

@njit
def init_status(problem):
    N, M, G, P = problem
    used = np.zeros(N, np.bool_)
    count = np.zeros(N, np.int64)
    for _ in range(M):
        while True:
            i = np.random.randint(0, N)
            if used[i]:
                continue
            break
        used[i] = 1
        for j in range(N):
            if G[i, j]:
                count[j] += 1
    score = 0
    for i in range(N):
        if count[i]:
            score += P[i]
    return (used, count), score

@njit
def calc_score(problem, state):
    N, M, G, P = problem
    used, count = state
    score = 0
    for i in range(N):
        if count[i]:
            score += P[i]
    return score

@njit
def move(problem, state, score, param):
    N, M, G, P = problem
    used, count = state
    i, j = param
    change_log = (i, j)
    used[i] = 0
    used[j] = 1
    for v in range(N):
        if count[v]:
            score -= P[v]
        count[v] = count[v] - G[i, v] + G[j, v]
        if count[v]:
            score += P[v]
    return (used, count), score, change_log


@njit
def rollback(problem, state, score, change_log):
    i, j = change_log
    state, score, change_log = move(problem, state, score, (j, i))
    return state, score

@njit
def mountain(problem, state, score, max_step):
    N, M, G, P = problem
    assert calc_score(problem, state) == score

    def random_param():
        used, count = state
        while True:
            i = np.random.randint(0, N)
            if used[i]:
                break
        while True:
            j = np.random.randint(0, N)
            if not used[j]:
                break
        return i, j

    for step in range(max_step):
        param = random_param()
        before_score = score
        state, score, change_log = move(problem, state, score, param)
        if before_score > score:
            state, score = rollback(problem, state, score, change_log)
    return state, score

@njit
def get_graph(N, XY, D):
    G = np.zeros((N, N), np.bool_)
    for i in range(N):
        for j in range(N):
            dx, dy = XY[i] - XY[j]
            G[i, j] = dx * dx + dy * dy <= D * D
    return G

def output(state, score, file=None):
    if file is None:
        file = sys.stdout
    else:
        file = open(file, 'w')
    used, count = state
    A = np.where(used)[0]
    for v in A:
        print(v + 1, file=file)
    print(score, file=file)

def main(file=None):
    if file:
        sys.stdin = open(file)
    read = sys.stdin.buffer.read
    readline = sys.stdin.buffer.readline
    readlines = sys.stdin.buffer.readlines

    N = int(readline())
    M = int(readline())
    D = int(readline())
    XYP = np.array(read().split(), np.int64).reshape(N, 3)
    XY = XYP[:, :2]
    G = get_graph(N, XY, D)
    P = XYP[:, 2].copy()
    problem = (N, M, G, P)
    while True:
        state, score = init_status(problem)
        while True:
            before_score = score
            state, score = mountain(problem, state, score, 10**5)
            if score == before_score:
                break
            f_name = f'{file}_{score}_out.txt'
            output(state, score, f_name)

main('test/sample-5.in')"""
print("""3
6
11
15
18
78100""")

Submission Info

Submission Time
Task election1 - 選挙 (Election)
User maspy
Language Python (3.8.2)
Score 20
Code Size 3456 Byte
Status AC
Exec Time 19 ms
Memory 8780 KiB

Judge Result

Set Name Case01
Score / Max Score 20 / 21
Status
AC × 1
Set Name Test Cases
Case01 01
Case Name Status Exec Time Memory
01 AC 19 ms 8780 KiB