'''runda 2B: stos naleśników'''
import sys
from heapq import heapify, heappop, heappush
from io import StringIO
from collections import namedtuple
from itertools import chain
def to_ints(line: str) -> list[int]:
return [int(x) for x in line.split()]
def to_lists_of_ints(s: str) -> list[list[int]]:
return [to_ints(line) for line in s.strip().split('\n')]
DEBUG = False
def debug(*args): print(*args) if DEBUG else None
input = '''
3 3 5
1 2 3
1 2 3
3 2 1
'''
input = '''
8 2 3
255 256
687 717
242 248
576 586
19 48
735 711
145 101
554 558
''' # expected 2139
input = '''
3 2 3
633 619
786 825
933 904
''' # exp 2623
debug(input)
instream = StringIO(input.strip())
instream = sys.stdin
n, stack_size, can_eat = to_ints(instream.readline())
stacks = to_lists_of_ints(instream.read())
assert n == len(stacks)
assert stack_size == len(stacks[0])
revs = [s for s in stacks if s[0] >= s[-1]]
norm = [s for s in stacks if s[0] < s[-1]]
best_from_revs = list(chain.from_iterable(revs))
best_from_revs.sort()
class SuperStack:
def __init__(self, nums):
# debug('init' + str(nums))
self.nums = sorted(nums)
for i in range(1, len(self.nums)):
self.nums[i] += self.nums[i-1]
# debug('init' + str(nums))
def __getitem__(self, i):
return self.nums[i] - (self.nums[i-1] if i != 0 else 0)
def __repr__(self):
nums = [x for x in self.nums]
for i in reversed(range(1, len(nums))):
nums[i] -= nums[i-1]
return str(nums)
def __len__(self):
return len(self.nums)
def __bool__(self):
return bool(self.nums)
def pop(self):
last = self.nums.pop()
return last - self.nums[-1] if self.nums else last
def avg_top(self, cnt):
if not self.nums: return 0
if cnt >= len(self.nums): return self.nums[-1] / len(self.nums)
return (self.nums[-1] - self.nums[-1 - cnt]) / cnt
# best_from_revs = SuperStack(best_from_revs)
# best_from_revs = []
# heap = [(-s[0], list(reversed(s))) for s in revs]
# heapify(heap)
#
# while heap and len(best_from_revs) < can_eat:
# (prio, stack) = heappop(heap)
# best = stack.pop()
# best_from_revs.append(best)
# assert best == -prio
# if stack:
# heappush(heap, (-stack[-1], stack))
#
# best_from_revs.reverse()
# debug(best_from_revs)
stack = namedtuple('stack', 'prio suma list')
def new(l: list[int]):
suma = sum(l)
return stack(-suma/len(l), suma, l)
heap = [new(s) for s in norm]
heapify(heap)
################################################# old
result = 0
while can_eat:
if heap and can_eat < len(heap[0].list):
heap = [new(t.list[:can_eat]) for t in heap]
heapify(heap)
if not best_from_revs or (heap and -heap[0].prio >= best_from_revs[-1]):
tup = heappop(heap)
can_eat -= len(tup.list)
debug(tup.list)
result += tup.suma
else:
best = best_from_revs.pop()
debug(best)
result += best
can_eat -= 1
print(result)
#########################################333333333
# result = 0
# while can_eat:
# if heap and can_eat < len(heap[0].list):
# heap = [new(t.list[:can_eat]) for t in heap]
# heapify(heap)
# if not best_from_revs or (heap and -heap[0].prio >= best_from_revs.avg_top(len(heap[0].list)) and -heap[0].prio >= best_from_revs[-1]):
# avg = best_from_revs.avg_top(len(heap[0].list))
# tup = heappop(heap)
# can_eat -= len(tup.list)
# result += tup.suma
# debug("added:", tup.list, result, "cat_eat:", can_eat, 'avg:', avg, "superstack:", best_from_revs)
# else:
# best = best_from_revs.pop()
# result += best
# can_eat -= 1
# debug(f"added: {best}", "result:", result, "cat_eat:", can_eat)
#
# print(result)
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 | '''runda 2B: stos naleśników''' import sys from heapq import heapify, heappop, heappush from io import StringIO from collections import namedtuple from itertools import chain def to_ints(line: str) -> list[int]: return [int(x) for x in line.split()] def to_lists_of_ints(s: str) -> list[list[int]]: return [to_ints(line) for line in s.strip().split('\n')] DEBUG = False def debug(*args): print(*args) if DEBUG else None input = ''' 3 3 5 1 2 3 1 2 3 3 2 1 ''' input = ''' 8 2 3 255 256 687 717 242 248 576 586 19 48 735 711 145 101 554 558 ''' # expected 2139 input = ''' 3 2 3 633 619 786 825 933 904 ''' # exp 2623 debug(input) instream = StringIO(input.strip()) instream = sys.stdin n, stack_size, can_eat = to_ints(instream.readline()) stacks = to_lists_of_ints(instream.read()) assert n == len(stacks) assert stack_size == len(stacks[0]) revs = [s for s in stacks if s[0] >= s[-1]] norm = [s for s in stacks if s[0] < s[-1]] best_from_revs = list(chain.from_iterable(revs)) best_from_revs.sort() class SuperStack: def __init__(self, nums): # debug('init' + str(nums)) self.nums = sorted(nums) for i in range(1, len(self.nums)): self.nums[i] += self.nums[i-1] # debug('init' + str(nums)) def __getitem__(self, i): return self.nums[i] - (self.nums[i-1] if i != 0 else 0) def __repr__(self): nums = [x for x in self.nums] for i in reversed(range(1, len(nums))): nums[i] -= nums[i-1] return str(nums) def __len__(self): return len(self.nums) def __bool__(self): return bool(self.nums) def pop(self): last = self.nums.pop() return last - self.nums[-1] if self.nums else last def avg_top(self, cnt): if not self.nums: return 0 if cnt >= len(self.nums): return self.nums[-1] / len(self.nums) return (self.nums[-1] - self.nums[-1 - cnt]) / cnt # best_from_revs = SuperStack(best_from_revs) # best_from_revs = [] # heap = [(-s[0], list(reversed(s))) for s in revs] # heapify(heap) # # while heap and len(best_from_revs) < can_eat: # (prio, stack) = heappop(heap) # best = stack.pop() # best_from_revs.append(best) # assert best == -prio # if stack: # heappush(heap, (-stack[-1], stack)) # # best_from_revs.reverse() # debug(best_from_revs) stack = namedtuple('stack', 'prio suma list') def new(l: list[int]): suma = sum(l) return stack(-suma/len(l), suma, l) heap = [new(s) for s in norm] heapify(heap) ################################################# old result = 0 while can_eat: if heap and can_eat < len(heap[0].list): heap = [new(t.list[:can_eat]) for t in heap] heapify(heap) if not best_from_revs or (heap and -heap[0].prio >= best_from_revs[-1]): tup = heappop(heap) can_eat -= len(tup.list) debug(tup.list) result += tup.suma else: best = best_from_revs.pop() debug(best) result += best can_eat -= 1 print(result) #########################################333333333 # result = 0 # while can_eat: # if heap and can_eat < len(heap[0].list): # heap = [new(t.list[:can_eat]) for t in heap] # heapify(heap) # if not best_from_revs or (heap and -heap[0].prio >= best_from_revs.avg_top(len(heap[0].list)) and -heap[0].prio >= best_from_revs[-1]): # avg = best_from_revs.avg_top(len(heap[0].list)) # tup = heappop(heap) # can_eat -= len(tup.list) # result += tup.suma # debug("added:", tup.list, result, "cat_eat:", can_eat, 'avg:', avg, "superstack:", best_from_revs) # else: # best = best_from_revs.pop() # result += best # can_eat -= 1 # debug(f"added: {best}", "result:", result, "cat_eat:", can_eat) # # print(result) |
English