Complex Market Analysis (Codeforces 1609C)
Tags: Implement, schedules

  While performing complex market analysis William encountered the following problem:For a given array $$$a$$$ of size $$$n$$$ and a natural number $$$e$$$, calculate the number of pairs of natural numbers $$$(i, k)$$$ which satisfy the following conditions:   $$$1 \le i, k$$$  $$$i + e \cdot k \le n$$$.  Product $$$a_i \cdot a_{i + e} \cdot a_{i + 2 \cdot e} \cdot \ldots \cdot a_{i + k \cdot e} $$$ is a prime number. A prime number (or a prime) is a natural number greater than 1 that is not a product of two smaller natural numbers.

Input: Each test contains multiple test cases. The first line contains the number of test cases $$$t$$$ ($$$1 \le t \le 10\,000$$$). Description of the test cases follows.The first line of each test case contains two integers $$$n$$$ and $$$e$$$ $$$(1 \le e \le n \le 2 \cdot 10^5)$$$, the number of items in the array and number $$$e$$$, respectively.The second line contains $$$n$$$ integers $$$a_1, a_2, \dots, a_n$$$ $$$(1 \le a_i \le 10^6)$$$, the contents of the array.It is guaranteed that the sum of $$$n$$$ over all test cases does not exceed $$$2 \cdot 10^5$$$.

Output: For each test case output the answer in the following format:Output one line containing the number of pairs of numbers $$$(i, k)$$$ which satisfy the conditions.

Note: In the first example test case two pairs satisfy the conditions:   $$$i = 2, k = 1$$$, for which the product is: $$$a_{2} \cdot a_{5} = 2$$$ which is a prime number.  $$$i = 3, k = 1$$$, for which the product is: $$$a_{3} \cdot a_{6} = 19$$$ which is a prime number. In the second example test case there are no pairs that satisfy the conditions.In the third example test case four pairs satisfy the conditions:   $$$i = 1, k = 1$$$, for which the product is: $$$a_{1} \cdot a_{4} = 2$$$ which is a prime number.  $$$i = 1, k = 2$$$, for which the product is: $$$a_{1} \cdot a_{4} \cdot a_{7} = 2$$$ which is a prime number.  $$$i = 3, k = 1$$$, for which the product is: $$$a_{3} \cdot a_{6} = 2$$$ which is a prime number.  $$$i = 6, k = 1$$$, for which the product is: $$$a_{6} \cdot a_{9} = 2$$$ which is a prime number. In the fourth example test case there are no pairs that satisfy the conditions.In the fifth example test case five pairs satisfy the conditions:   $$$i = 1, k = 1$$$, for which the product is: $$$a_{1} \cdot a_{2} = 2$$$ which is a prime number.  $$$i = 1, k = 2$$$, for which the product is: $$$a_{1} \cdot a_{2} \cdot a_{3} = 2$$$ which is a prime number.  $$$i = 1, k = 3$$$, for which the product is: $$$a_{1} \cdot a_{2} \cdot a_{3} \cdot a_{4} = 2$$$ which is a prime number.  $$$i = 2, k = 1$$$, for which the product is: $$$a_{2} \cdot a_{3} = 2$$$ which is a prime number.  $$$i = 2, k = 2$$$, for which the product is: $$$a_{2} \cdot a_{3} \cdot a_{4} = 2$$$ which is a prime number. In the sixth example test case there are no pairs that satisfy the conditions.

Sample Inputs:
6
7 3
10 2 1 3 1 19 3
3 2
1 13 1
9 3
2 4 2 1 1 1 1 4 2
3 1
1 1 1
4 1
1 2 1 1
2 2
1 2

Sample Outputs:
2
0
4
0
5
0

Solution (Python):
import sys, os, io
from collections import *
MOD = int(1000000000.0) + 7
inf = float('inf')

def fileInput():
    sys.stdin = open('input.txt', 'r')

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

def integer():
    return int(input())

def string():
    return input()

def matrixNum(m, n):
    return [lst() for i in range(m)]

def cond(a, h, m):
    for i in a:
        if i > m:
            h -= i
        if h <= 0:
            return False
    return True

class DSU:

    def __init__(self, n) -> None:
        self.parent = [i for i in range(n)]
        self.rank = [0] * n

    def find(self, x):
        if self.parent[x] == x:
            return x
        return self.find(self.parent[x])

    def union(self, x, y):
        parentX = self.find(x)
        parentY = self.find(y)
        if parentX == parentY:
            return 0
        if self.rank[parentX] > self.rank[parentY]:
            self.parent[parentY] = parentX
        elif self.rank[parentX] < self.rank[parentY]:
            self.parent[parentX] = parentY
        else:
            self.parent[parentX] = parentY
            self.rank[parentY] += 1
        return 1
prime = [1] * int(1000000.0 + 5)
prime[0] = prime[1] = 0
n = len(prime)
for i in range(2, int(n ** 0.5) + 1):
    if prime[i]:
        for j in range(i * i, n, i):
            prime[j] = 0

def solve():
    (n, e) = lst()
    a = lst()
    cnt = 0
    for i in range(n):
        if prime[a[i]]:
            left = 0
            right = 1
            for j in range(i - e, -1, -e):
                if a[j] == 1:
                    left += 1
                else:
                    break
            for j in range(i + e, n, e):
                if a[j] == 1:
                    right += 1
                else:
                    break
            cnt += left * right + (right - 1)
    print(cnt)
t = 1
t = integer()
for _ in range(t):
    solve()

Solution (Python):
from __future__ import annotations
from sys import stdin
import os
from math import isqrt
dbg = os.getenv('DEBUG')

def make_sieve(n: int) -> bytearray:
    sieve = bytearray([1] * (n + 1))
    sieve[0] = sieve[1] = 0
    for x in range(2, isqrt(n) + 1):
        if not sieve[x]:
            continue
        for y in range(2 * x, n + 1, x):
            sieve[y] = 0
    return sieve
SIEVE = make_sieve(10 ** 6)

def solve(ls: list[int]) -> int:
    ys: list[list[int]] = []
    ys.append([ls[0], 1])
    for x in ls[1:]:
        if x == 1 and ys[-1][0] == 1:
            ys[-1][1] += 1
            continue
        ys.append([x, 1])
    result = 0
    k = len(ys)
    for (i, [x, _]) in enumerate(ys):
        if not SIEVE[x]:
            continue
        left = 1
        right = 1
        if i > 0:
            if ys[i - 1][0] == 1:
                left += ys[i - 1][1]
        if i < k - 1:
            if ys[i + 1][0] == 1:
                right += ys[i + 1][1]
        result += left * right - 1
    return result

def main_case() -> None:
    (_, e) = map(int, stdin.readline().split())
    ls = list(map(int, stdin.readline().split()))
    result = sum((solve(ls[i::e]) for i in range(e)))
    print(result)

def main() -> None:
    t = int(stdin.readline())
    for _ in range(t):
        main_case()
if __name__ == '__main__':
    main()
''

Solution (Python):
from collections import defaultdict
import sys
p2D = lambda x: print(*x, sep='\n')

def II():
    return int(sys.stdin.buffer.readline())

def MI():
    return map(int, sys.stdin.buffer.readline().split())

def LI():
    return list(map(int, sys.stdin.buffer.readline().split()))

def LLI(rows_number):
    return [LI() for _ in range(rows_number)]

def BI():
    return sys.stdin.buffer.readline().rstrip()

def SI():
    return sys.stdin.buffer.readline().rstrip().decode()

def li():
    return [int(i) for i in input().split()]

def lli(rows):
    return [li() for _ in range(rows)]

def si():
    return input()

def ii():
    return int(input())

def ins():
    return input().split()

def drank(d, processing, da, rank):
    tmp = 10 ** 9
    if len(d[da]) == 1:
        return 1
    for di in d[da]:
        if processing[di - 1] == 0:
            processing[di - 1] = 1
            tmp = min(tmp, drank(d, processing, di, rank))
            processing[di - 1] = 0
    rank[da - 1] = tmp + 1
    return tmp + 1
primes = []

def primes_till_n(n):
    composite = [0] * (n + 1)
    for a in range(2, n + 1):
        i = 2
        if composite[a] == 0:
            primes.append(a)
            while a * i <= n:
                composite[a * i] = 1
                i += 1

def binary_search(n, a):
    l = len(a)
    low = 0
    high = l - 1
    while high >= low:
        mid = (high + low) // 2
        if a[mid] == n:
            return mid
        elif a[mid] > n:
            high = mid - 1
        else:
            low = mid + 1
    return -1
is_prime = [True] * (10 ** 6 + 1)

def primeseive(n):
    is_prime[0] = is_prime[1] = False
    i = 2
    while i * i <= n:
        if is_prime[i] == True:
            for j in range(i * i, n + 1, i):
                is_prime[j] = False
        i += 1

def solve():
    (n, e) = LI()
    a = LI()
    ans = 0
    for i in range(n):
        if is_prime[a[i]] == True:
            j = i
            x = 0
            y = 0
            while j + e < n and a[j + e] == 1:
                x += 1
                j += e
            j = i
            while j - e > -1 and a[j - e] == 1:
                y += 1
                j -= e
            ans += x + y + x * y
    return ans

def main():
    primeseive(10 ** 6)
    for _ in range(II()):
        sys.stdout.write(str(solve()) + '\n')
if __name__ == '__main__':
    main()