【校内模拟赛】Day2-A

链接

D2-A

题目描述

n个小球排成一列，第i个小球颜色为ai,定义小球i与j的距离为|(i-j)*(ai-aj)|,求所有小球间距离之和

样例输入

5
0 1 1 0 1

样例输出

11

思路

O(n)算法
当前位为1时，只有前面为0的时候才能产生贡献，当前位为0的时候就相反，那么就记录0的位置和和个数，1的位置和和个数
直接读入的时候算就好了

代码

#include<iostream>
#include<cstdio>
#include<cstring>
#include<algorithm>
#define ll long long

using namespace std;

ll n, x, ans, num0, num1, num0_sum, num1_sum;

ll read() 
{
	ll x = 0, flag = 1; char ch = getchar();
	while (ch < '0' || ch > '9') {if (ch == '-') flag = -1; ch = getchar();}
	while (ch >= '0' && ch <= '9'){x = x * 10 + ch - '0'; ch = getchar();}
	return x * flag;
}

int main()
{
	n = read();
	for(int i = 1; i <= n; ++i)	
	{
		x = read();
		if(x == 1)
		{
			ans += num0 * i - num0_sum;
			num1++; num1_sum += i;
		}
		else {
 			ans += num1 * i - num1_sum;
			num0++; num0_sum += i;
		}
	}
	printf("%lld", ans);
	return 0;
}