Codeforces Round #133 (Div. 2) E : Martian Luck

問題概要

K(<10^9)進数でかかれたN(<10^5)桁の数字が(実際には配列として)与えられる。digital rootがBになるような部分文字列(leading zeroも可)がいくつあるか求める問題。

解法

digital rootはmod (K-1)で考えればよい(0に注意)。後はmapに累積和を詰め込みながらカウントすればよい。B=0やB=K-1のときにdigital rootが0の場合とK-1の場合をちゃんと区別するようにする。

acceptされたコード

#include <cstdio>
#include <map>
using namespace std;

typedef long long int64;

const int MAX_N = int(1e5);

int K, B, N;
int as[MAX_N];

void init() {
	scanf("%d%d%d", &K, &B, &N);
	for (int i = 0; i < N; ++i) {
		scanf("%d", as + i);
	}
}

int64 solve() {
	int64 ans = 0;
	int64 countZero = 0;
	for (int i = 0, cnt = 0; i < N; ++i) {
		if (as[i]) {
			cnt = 0;
		}
		else {
			++cnt;
		}
		countZero += cnt;
	}
	if (B == 0) {
		return countZero;
	}
	map<int64, int> dict;
	++dict[0];

	int acc = 0;
	for (int i = 0; i < N; ++i) {
		acc = (acc + as[i]) % (K - 1);
		int s = (acc - B + (K - 1)) % (K - 1);
		ans += dict[s];
		++dict[acc];
	}
	if (B == K - 1) {
		ans -= countZero;
	}

	return ans;
}

int main() {
	init();
	printf("%lld\n", solve());
	return 0;
}