4128 - Iluminat

Primarul orașului X dorește să aibă un iluminat public modern. Pentru aceasta, realizează o schiță sub forma unui pătrat cu n linii și n coloane în care fiecare element situat la intersecția unei linii cu o coloană reprezintă un cartier.

Primarul a calculat pentru fiecare cartier care este numărul de stâlpi de iluminat public din acel cartier. Fiecare stâlp are un singur bec care inițial este aprins. Acesta a observat un lucru interesant: toate cartierele au un număr diferit de stâlpi de iluminare, iar valoarea maximă a numărului de stâlpi dintr-un cartier este n2.

Pentru a fi realizată într-un mod cât mai eficient, stingerea becurilor se realizează în următorul mod:

• în prima etapă se sting becurile din cartierul cu număr maxim de stâlpi de iluminat, ceea ce duce la stingerea becurilor din cartierele de pe aceeași linie precum și din cele pe aceeași coloană cu cartierul cu număr maxim de stâlpi.
• procedeul se reia la fiecare etapă pentru toate cartierele în care nu au fost stinse becurile, până când rămâne un singur cartier iluminat.

Cerința

Cunoscând numerele naturale nenule n și k, precum și numărul de stâlpi de iluminat din fiecare cartier, să se determine:

• Câți stâlpi de iluminat se află în cartierul cu număr maxim de stâlpi de iluminat la etapa cu numărul k din procedeul de stingere a becurilor?
• Câte becuri se sting, în total, la etapa cu numărul k?
• Care este numărul maxim de becuri aprinse într-o zonă pătratică a orașului de dimensiune k x k, înainte de a începe stingerea becurilor?

Date de intrare

Fișierul de intrare iluminat.in conține pe prima linie o cifră c (1, 2 sau 3), reprezentând cerința cerută. Pe linia următoare se găsesc două numere naturale nenule n și k, separate printr-un spațiu. Pe următoarele n linii se află n2 numere naturale distincte, câte n pe fiecare linie, separate prin câte un spațiu, cu semnificația din enunț.

Date de ieșire

În fișierul de ieșire iluminat.out se va afișa răspunsul în funcție de cerință:

• dacă c = 1 se va afișa pe prima linie un singur număr reprezentând numărul de stâlpi de iluminat din cartierul cu număr maxim de stâlpi de iluminat la etapa k
• dacă c = 2 se va afișa pe prima linie un singur număr reprezentând câte becuri se sting, în total, la etapa cu numărul k
• dacă c = 3 se va afișa numărul maxim de becuri aprinse într-o zonă pătratică de dimensiune k x k înainte de stingerea becurilor

Restricții și precizări

• c ∈ {1, 2, 3}
• 1 ≤ k < n ≤ 1000
• Numărul de becuri din fiecare cartier este mai mic sau egal cu n2

Exemplu:

iluminat.in

1
4 2
1 2 3 4
16 13 5 6
12 9 7 14
10 11 8 15

iluminat.out

15

Explicație

Cerința este 1. Se sting becurile din cartierul având 16 stâlpi de iluminat, ceea ce duce la stingerea becurilor de pe stâlpii din linia 2 și din coloana 1. Tabloul devine:

0 2 3 4

0 0 0 0

0 9 7 14

0 11 8 15

La etapa a doua, primul cartier în care se sting becurile are 15 stâlpi de iluminat.

import numpy as np

def generate_matrix(n):
    return np.arange(1, n*n + 1).reshape(n, n)

def find_max_and_sting(matrix):
    max_value = np.max(matrix)
    max_index = np.unravel_index(np.argmax(matrix, axis=None), matrix.shape)
    row, col = max_index
    matrix[row, :] = 0
    matrix[:, col] = 0
    return max_value, np.count_nonzero(matrix == 0)

def max_bulbs_in_kxk(matrix, k):
    max_sum = 0
    for i in range(matrix.shape[0] - k + 1):
        for j in range(matrix.shape[1] - k + 1):
            submatrix_sum = np.sum(matrix[i:i+k, j:j+k])
            if submatrix_sum > max_sum:
                max_sum = submatrix_sum
    return max_sum

def stage_results(n, k):
    matrix = generate_matrix(n)
    original_matrix = matrix.copy()
    
    max_value_at_k = 0
    total_sting_at_k = 0
    
    for stage in range(1, k + 1):
        max_value, sting_count = find_max_and_sting(matrix)
        if stage == k:
            max_value_at_k = max_value
            total_sting_at_k = sting_count

    max_in_kxk_before_sting = max_bulbs_in_kxk(original_matrix, k)
    
    return max_value_at_k, total_sting_at_k, max_in_kxk_before_sting

# Citirea datelor de intrare
n = int(input())
k = int(input())

# Calcularea și afișarea rezultatelor
max_value_at_k, total_sting_at_k, max_in_kxk_before_sting = stage_results(n, k)

print(max_value_at_k)
print(total_sting_at_k)
print(max_in_kxk_before_sting)