2779 - Cnt SQ

Cerința

Se dă o matrice binară (valori 0 și 1). Să se determine câte pătrate există cu proprietatea că acestea au pe marginea lor doar valori 1.

Date de intrare

Fișierul de intrare cntsq.in conține pe prima linie numerele N și M, reprezentând numărul de linii și numărul de coloane ale matricei, apoi N linii, pe fiecare câte M valori 0 sau 1, neseparate prin niciun spațiu, reprezentând elementele matricei.

Date de ieșire

Fișierul de ieșire cntsq.out va conține pe prima linie numărul C, reprezentând numărul de pătrate ce conțin doar 1 pe marginea lor.

Restricții și precizări

• 1 ≤ N, M ≤ 1000
• Se consideră pătrat și cel de latură 1 (conține doar un element).
• Fie un pătrat determinat de (i1, j1) colțul stânga sus, (i2, j2) colțul dreapta jos. Se definește marginea pătratului ca fiind mulțimea de puncte (x, y) care respectă cel puțin una din condițiile: ((x = i1 sau x = i2) și (j1 ≤ y ≤ j2)) sau ((y = j1 sau y = j2) și (i1 ≤ x ≤ i2)).

Exemplu:

cntsq.in

3 3
111
101
111

cntsq.out

9

cntsq.in

7 7
0000000
0111100
0101111
0100101
0111111
0000011
0000011

cntsq.out

27

Rezolvare

def count_squares(matrix):
    n, m = len(matrix), len(matrix[0])
    dp = [[0] * (m + 1) for _ in range(n + 1)]
    count = 0

    for i in range(1, n + 1):
        for j in range(1, m + 1):
            if matrix[i - 1][j - 1] == 1:
                dp[i][j] = min(dp[i - 1][j - 1], dp[i - 1][j], dp[i][j - 1]) + 1
                size = dp[i][j]
                while size > 0:
                    count += 1
                    size -= 1

    return count

def main():
    with open("cntsq.in", "r") as fin, open("cntsq.out", "w") as fout:
        n, m = map(int, fin.readline().split())
        matrix = []
        for _ in range(n):
            row = list(map(int, fin.readline().strip()))
            matrix.append(row)

        result = count_squares(matrix)
        fout.write(str(result) + "\n")

if __name__ == "__main__":
    main()