M72NRSH-Editorial

general
,
1

PROBLEM LINK:

Practice
Contest

Author: manjunath1996

Editorialist: manjunath1996

DIFFICULTY:

EASY-MEDIUM

PREREQUISITES:

DP

EXPLANATION:

The brute force approach is to check for all j< i where j-i=cK and a[j]=a[i].In the worst case, it takes O(n^2) time which will give TLE.
We should note that a[i] is in between 1 to 100. So a better approach will be to keep a 2D dp state of index and element . dp[i][x] will store number of indices j<=i where a[j]=a[i] and j-i=c*K for c>=0;

now we our answer will ( \sum_{i=0}^{n-1}\sum_{j=1}^{100}(dp[i][j]) )-n .n is substracted because we should exclude cases where c=0.

long long ans=0;
for i=0 to n-1 {
	for j=1 to 100 {
            if(i>=K)
		dp[i][j]=dp[i-K][j]
	}
        if(i>=K)
	     ans+=dp[i-K][a[i]]
	dp[i][a[i]]++;
}

so, the final answer can be calculated with O(Z*n) where Z<=100. So it passess given time limit.

Note we should use long long for an answer as it may exceed integer range.

Also, there is an improved approach which solves this question in O(n) time complexity.

long long ans=0;
for i=0 to n-1 {
        ans+=dp[i%K][a[i]]
	dp[i%K][a[i]]++;
}

AUTHOR’S AND TESTER’S SOLUTIONS:

Author’s solution can be found here.