Zbazi in Zeydabad (Codeforces 628E)
Tags: Implement, data structures

A tourist wants to visit country Zeydabad for Zbazi (a local game in Zeydabad).The country Zeydabad is a rectangular table consisting of n rows and m columns. Each cell on the country is either 'z' or '.'.The tourist knows this country is named Zeydabad because there are lots of ''Z-pattern"s in the country. A ''Z-pattern" is a square which anti-diagonal is completely filled with 'z' and its upper and lower rows are also completely filled with 'z'. All other cells of a square can be arbitrary.  Note that a ''Z-pattern" can consist of only one cell (see the examples).So he wants to count the number of ''Z-pattern"s in the country (a necessary skill for Zbazi).Now your task is to help tourist with counting number of ''Z-pattern"s.As input/output can reach huge size it is recommended to use fast input/output methods: for example, prefer to use gets/scanf/printf instead of getline/cin/cout in C++, prefer to use BufferedReader/PrintWriter instead of Scanner/System.out in Java.

Input: The first line contains two integers n, m (1 ≤ n, m ≤ 3000) — the number of rows and columns respectively.Each of the next n lines contains m characters 'z' or '.' — the description of Zeydabad.

Output: Print the only integer a — the number of ''Z-pattern"s in Zeydabad.


Sample Inputs:
4 4
zzzz
zzz.
.z..
zzzz

1 4
z.z.

2 2
zz
zz

Sample Outputs:
16

2

5

Solution (Python):
import sys, os, io
input = io.BytesIO(os.read(0, os.fstat(0).st_size)).readline

def make_tree(n):
    tree1 = [0] * (n + 1)
    tree2 = [0] * (n + 1)
    return (tree1, tree2)

def add(i, x, tree):
    while i <= c:
        tree[i] += x
        i += i & -i

def add_segment(l, r, x):
    add(l, -x * (l - 1), tree1)
    add(r + 1, x * r, tree1)
    add(l, x, tree2)
    add(r + 1, -x, tree2)

def get_sum(i, tree):
    s = 0
    while i > 0:
        s += tree[i]
        i -= i & -i
    return s

def get_sum_segment(s, t):
    s -= 1
    x = get_sum(s, tree1) + get_sum(s, tree2) * s
    y = get_sum(t, tree1) + get_sum(t, tree2) * t
    return y - x
(n, m) = map(int, input().split())
l = [[] for _ in range(n + m)]
r = [[] for _ in range(n + m)]
for i in range(n):
    s = list(input().rstrip())
    c = 0
    for j in range(m):
        if s[j] == 122:
            c += 1
        else:
            c = 0
        l[i + j].append(c)
    c = 0
    for j in range(m - 1, -1, -1):
        if s[j] == 122:
            c += 1
        else:
            c = 0
        r[i + j].append(c)
for i in range(n + m):
    l[i].append(0)
    r[i].append(0)
(ans1, ans2) = (0, 0)
mod = pow(10, 9) + 7
for i in range(n + m):
    (l0, r0) = (l[i], r[i])
    (li, ri) = ([], [])
    for (u, v) in zip(l0, r0):
        if not u + v and li:
            c = len(li)
            x = [[] for _ in range(c + 1)]
            k = 0
            for j in range(c - 1, -1, -1):
                x[max(0, c - ri[j] - k)].append(c - j)
                k += 1
            (tree1, tree2) = make_tree(c)
            for j in range(c - 1, -1, -1):
                t = c - j
                s = max(1, t - li[j] + 1)
                add_segment(s, t, 1)
                if x[j]:
                    (s, t) = (-1, -1)
                    for k in x[j]:
                        if s == -1:
                            (s, t) = (k, k)
                        elif t + 1 < k:
                            ans1 += get_sum_segment(s, t)
                            (s, t) = (k, k)
                        else:
                            t += 1
                    ans1 += get_sum_segment(s, t)
            ans2 += ans1 // mod
            ans1 %= mod
            (li, ri) = ([], [])
        elif u + v:
            li.append(u)
            ri.append(v)
ans = ans1 + ans2 * mod
print(ans)

Solution (Python):
(standard_input, packages, output_together) = (1, 1, 0)
(dfs, hashing, read_from_file) = (0, 0, 0)
de = 1
if 1:
    if standard_input:
        import io, os, sys
        input = lambda : sys.stdin.readline().strip()
        import math
        inf = math.inf

        def I():
            return input()

        def II():
            return int(input())

        def MII():
            return map(int, input().split())

        def LI():
            return list(input().split())

        def LII():
            return list(map(int, input().split()))

        def LFI():
            return list(map(float, input().split()))

        def GMI():
            return map(lambda x: int(x) - 1, input().split())

        def LGMI():
            return list(map(lambda x: int(x) - 1, input().split()))
    if packages:
        from io import BytesIO, IOBase
        import random
        import os
        import bisect
        import typing
        from collections import Counter, defaultdict, deque
        from copy import deepcopy
        from functools import cmp_to_key, reduce, lru_cache
        from heapq import merge, heapify, heappop, heappush, heappushpop, nlargest, nsmallest
        from itertools import accumulate, combinations, permutations, count, product
        from operator import add, iand, ior, itemgetter, mul, xor
        from string import ascii_lowercase, ascii_uppercase, ascii_letters
        from typing import *
        BUFSIZE = 4096
    if output_together:

        class FastIO(IOBase):
            newlines = 0

            def __init__(self, file):
                self._fd = file.fileno()
                self.buffer = BytesIO()
                self.writable = 'x' in file.mode or 'r' not in file.mode
                self.write = self.buffer.write if self.writable else None

            def read(self):
                while True:
                    b = os.read(self._fd, max(os.fstat(self._fd).st_size, BUFSIZE))
                    if not b:
                        break
                    ptr = self.buffer.tell()
                    (self.buffer.seek(0, 2), self.buffer.write(b), self.buffer.seek(ptr))
                self.newlines = 0
                return self.buffer.read()

            def readline(self):
                while self.newlines == 0:
                    b = os.read(self._fd, max(os.fstat(self._fd).st_size, BUFSIZE))
                    self.newlines = b.count(b'\n') + (not b)
                    ptr = self.buffer.tell()
                    (self.buffer.seek(0, 2), self.buffer.write(b), self.buffer.seek(ptr))
                self.newlines -= 1
                return self.buffer.readline()

            def flush(self):
                if self.writable:
                    os.write(self._fd, self.buffer.getvalue())
                    (self.buffer.truncate(0), self.buffer.seek(0))

        class IOWrapper(IOBase):

            def __init__(self, file):
                self.buffer = FastIO(file)
                self.flush = self.buffer.flush
                self.writable = self.buffer.writable
                self.write = lambda s: self.buffer.write(s.encode('ascii'))
                self.read = lambda : self.buffer.read().decode('ascii')
                self.readline = lambda : self.buffer.readline().decode('ascii')
        sys.stdout = IOWrapper(sys.stdout)
    if dfs:
        from types import GeneratorType

        def bootstrap(f, stack=[]):

            def wrappedfunc(*args, **kwargs):
                if stack:
                    return f(*args, **kwargs)
                else:
                    to = f(*args, **kwargs)
                    while True:
                        if type(to) is GeneratorType:
                            stack.append(to)
                            to = next(to)
                        else:
                            stack.pop()
                            if not stack:
                                break
                            to = stack[-1].send(to)
                    return to
            return wrappedfunc
    if hashing:
        RANDOM = random.getrandbits(20)

        class Wrapper(int):

            def __init__(self, x):
                int.__init__(x)

            def __hash__(self):
                return super(Wrapper, self).__hash__() ^ RANDOM
    if read_from_file:
        file = open('input.txt', 'r').readline().strip()[1:-1]
        fin = open(file, 'r')
        input = lambda : fin.readline().strip()
        output_file = open('output.txt', 'w')

        def fprint(*args, **kwargs):
            print(*args, **kwargs, file=output_file)
    if de:

        def debug(*args, **kwargs):
            print('\x1b[92m', end='')
            print(*args, **kwargs)
            print('\x1b[0m', end='')
(n, m) = MII()
strs = [I() for _ in range(n)]
lines = [0] * n
diag_len = [0] * n

class FenwickTree:

    def __init__(self, n):
        self.n = n
        self.bit = [0] * n

    def init(self):
        for i in range(self.n):
            self.bit[i] = 0

    def sum(self, r):
        res = 0
        while r >= 0:
            res += self.bit[r]
            r = (r & r + 1) - 1
        return res

    def rsum(self, l, r):
        return self.sum(r) - self.sum(l - 1)

    def add(self, idx, delta):
        while idx < self.n:
            self.bit[idx] += delta
            idx = idx | idx + 1
fen = FenwickTree(n)
ans = 0
st_range = list(range(n))
for j in range(m):
    for i in range(n - 1):
        diag_len[i] = diag_len[i + 1] + 1 if strs[i][j] == 'z' else 0
    diag_len[n - 1] = 1 if strs[n - 1][j] == 'z' else 0
    for i in range(n):
        lines[i] = lines[i] + 1 if strs[i][j] == 'z' else 0
    st_range.sort(key=lambda x: x - lines[x])
    pt = 0
    fen.init()
    for i in range(n):
        while pt < n and 1 + st_range[pt] - lines[st_range[pt]] <= i:
            fen.add(st_range[pt], 1)
            pt += 1
        v = diag_len[i] if diag_len[i] < lines[i] else lines[i]
        if v > 0:
            ans += fen.rsum(i, i + v - 1)
print(ans)

Solution (Python):
import sys
ni = lambda : int(input())
na = lambda : list(map(int, input().split()))
yes = lambda : print('yes')
Yes = lambda : print('Yes')
YES = lambda : print('YES')
no = lambda : print('no')
No = lambda : print('No')
NO = lambda : print('NO')

class BIT:

    def __init__(self, n):
        self.n = n
        self.data = [0] * (n + 1)
        self.el = [0] * (n + 1)

    def sum(self, i):
        s = 0
        while i > 0:
            s += self.data[i]
            i -= i & -i
        return s

    def add(self, i, x):
        assert 0 <= i < self.n
        i = i + 1
        self.el[i] += x
        while i <= self.n:
            self.data[i] += x
            i += i & -i

    def set(self, i, x):
        assert 0 <= i <= self.n
        self.add(i, x - self.get(i))

    def get(self, i, j=None):
        if j is None:
            return self.el[i + 1]
        return self.sum(j) - self.sum(i)
from collections import defaultdict

def solve(f):
    n = len(f)
    if len(f) == 0:
        return 0
    (a, b) = zip(*f)
    d = defaultdict(list)
    bit = BIT(n)
    for i in range(n):
        d[a[i] + i].append(i)
    for i in range(n):
        bit.add(i, 1)
    res = 0
    for i in range(n):
        for j in d[i]:
            bit.add(j, -1)
        res += bit.get(i - b[i] + 1, i + 1)
    return res
(h, w) = na()
c = [input() for i in range(h)]
''
m = min(h, w)
ans = 0
f = [[w + 1 for i in range(w)] for j in range(h)]
for i in range(h):
    if c[i][0] == 'z':
        f[i][0] = 0
for i in range(h):
    for j in range(1, w):
        if c[i][j] == 'z':
            f[i][j] = min(j, f[i][j - 1])
g = [[-1 for i in range(w)] for j in range(h)]
for i in range(h):
    if c[i][-1] == 'z':
        g[i][-1] = w - 1
for i in range(h):
    for j in range(w - 2, -1, -1):
        if c[i][j] == 'z':
            g[i][j] = max(j, g[i][j + 1])
for t in range(h + w - 1):
    sy = min(t, w - 1)
    sx = t - sy
    (x, y) = (sx, sy)
    i = 0
    l = 0
    ff = []
    while sx < h and sy >= 0:
        if c[sx][sy] == 'z':
            ff.append((sy - f[sx][sy] + 1, g[sx][sy] - sy + 1))
        else:
            ans += solve(ff)
            ff = []
        sx += 1
        sy -= 1
        i += 1
    ans += solve(ff)
print(ans)