Kostomuksha and AESC MSU Codechef Solution

by

Kostomuksha and AESC MSU Codechef Solution

A famous student of AESC MSU, as you know, comes from Kostomuksha. Kostomuksha has a popular game called Doka.

The essence of Doka is as follows:

You are given an array A and an integer X. You want to calculate how many subarrays of this array have a geometric mean of X.

Formally, calculate the number of pairs of integers (L,R) such that 1≤L≤R≤N and
AL⋅AL+1⋅…⋅AR−−−−−−−−−−−−−−−√R−L+1=X
Input Format
The first line of input contains an integer T, denoting the number of test cases. The description of T test cases follows.
Each test case consists of two lines of input.
The first line of each test case contains two space-separated integers N,X — the size of the array A and the required geometric mean.
The second line of each test case contains N space-separated integers A1,A2,…,AN.
Output Format
For each test case, on a new line print the number of subarrays that satisfy the condition.

Constraints
1≤T≤103
1≤N≤2⋅105
1≤X,Ai≤109
Sum of N over all test cases does not exceed 2⋅105.
Subtasks
Subtask #1 (100 points): Original constraints

Sample Input 1
3
3 3
3 3 3
4 4
1 2 3 4
4 54
36 81 54 54
Sample Output 1
6
1
6
Explanation
Test case 1: Every subarray has a geometric mean of 3, and there are 6 subarrays in total so the answer is 6.

Test case 2: The only subarray with a geometric mean of 4 is the singleton array [4], obtained with L=R=4.

Test case 3: The 6 pairs of (L,R) which have a geometric mean of 54 are:

(1,2), giving the subarray [36,81]
(1,3), giving the subarray [36,81,54]
(1,4), giving the subarray [36,81,54,54]
(3,3), giving the subarray [54]
(3,4), giving the subarray [54,54]
(4,4), giving the subarray [54]

SOLUTION

Program: Kostomuksha and AESC MSU Codechef Solution in Python

import numpy as np
for _ in range(int(input())):
    n,k=map(int,input().split())
    a=list(map(int,input().split()))
    res=[]
    for i in range(len(a)+1):
        for j in range(i):
            res.append(a[j:i])
    fl=[]
    for i in res:
        ans=np.log(i)
        fl.append(int(np.exp(ans.mean())))
    print(fl.count(k))

Program: Kostomuksha and AESC MSU Codechef Solution in C++

#include<bits/stdc++.h>
using namespace std;
#define Nx 35000
#define Mx 200050
bool vis[Nx];
int tcase, n, x;
int a[Mx];
int d[15];
vector<int> prime;
struct num {
    int l[15];
    num() {
        memset(l, 0, sizeof(l));
    }
} A, X;
map<num, int>mmp;
bool operator < (num A, num B) {
    for (int i = 0; i < 15; i++)
        if (A.l[i] != B.l[i]) return A.l[i] < B.l[i];
    return false;
}
num operator - (num A, num B) {
    for (int i = 0; i < 15; i++)
        A.l[i] -= B.l[i];
    return A;
}
void rds() {
    for (int i = 2; i < 35000; i++) {
        if (vis[i]) continue;
        prime.push_back(i);
        for (int j = i * i; j < 35000; j += i)
            vis[j] = 1;
    }
}
void ryd(int x) {
    for (int i = 0; i < 15; i++)
        d[i] = 0;
    int j = 0;
    for (int i = 0; i < prime.size(); i++) {
        if (x == 1) break;
        if (x % prime[i] == 0) {
            d[j] = prime[i];
            while (x % prime[i] == 0) {
                X.l[j]++;
                x /= prime[i];
            }
            j ++;
        }
    }
    if (x != 1) d[j] = x, X.l[j] = 1;
}
void kkk(int x) {
    for (int i = 0; i < 15; i++) {
        if (!d[i]) break;
        while (x % d[i] == 0) {
            A.l[i]++;
            x /= d[i];
        }
    }
    if (x > 1) A.l[14] ++;
}
int main() {
    ios::sync_with_stdio(0);
    cin.tie(0), cout.tie(0);
    rds();
    for (cin >> tcase; tcase--;) {
        cin >> n >> x;
        memset(A.l, 0, sizeof(A.l));
        memset(X.l, 0, sizeof(X.l));
        mmp.clear();
        ryd(x);
        mmp[A] = 1;
        for (int i = 0; i < n; i++) {
            cin >> a[i];
            kkk(a[i]);
            A = A - X;
            mmp[A] ++;
        }
        long long ans = 0;
        for (auto it = mmp.begin(); it != mmp.end(); it ++)
            ans += 1ll * (it -> second) * (it -> second - 1) / 2;
        cout << ans << endl;
    }
}

Related:

Leave a Comment

three × 2 =