Atcoder解いてみた AtCoder Beginner Contest C - Long Sequence

AtCoder Beginner Contest C - Long Sequence

問題はこちら

回答

最初の回答

途中まで解説動画で説明されていた通りに考えることができた
一応以下でもACなのだが、微妙な点がある

• C[ai]の回数分ループしていること
  Cは最大10^9なため、これだけループすると、TLEになりやすい
• if (tmpK_ - sub <= 0), if (tmpK__ - sub <= 0), if (K__ == 0)と分岐が多く、バグりやすい

#include <iostream>
#include <string>
#include <map>
#include <unordered_map>
#include <cstring>
#include <cmath>
#include <algorithm>
#include <vector>
#include <queue>
#include <stack>
#include <limits.h>
#include <bitset>
#include <list>
#include <set>
#include <numeric>
#include <tuple>

int N; // 1~
std::vector<int> L; // Li : 整数列Aiの長さ
std::vector<std::vector<int>> A; // 整数列 各値は、1~10^9
long long K; // 1~ N 10^9 2x10^5
std::vector<int> C;
// Liの合計値は、2x10^5までなため、整数列A内のすべての要素を見ることはできる

// Bの末尾に、Aiを連結させる操作をCi回 i=1,2,...Nの順
// Bの長さは最大ΣCiLi
std::vector<int> B;

int main()
{
	std::cin.tie(0);
	std::ios::sync_with_stdio(false);

	std::cin >> N >> K;
	A.resize(N);
	L.resize(N);
	C.resize(N);

	for (int i = 0; i < N; i++) {
		std::cin >> L[i];
		A[i].resize(L[i]);
		for (int l = 0; l < L[i]; l++) {
			int a;
			std::cin >> a;
			A[i][l] = a;
		}
	}

	for (int i = 0; i < N; i++) {
		std::cin >> C[i];
	}

	// KからLiCiを引いていく
	int ai;
	long long K_ = K;
	for (int i = 0; i < N; i++) {
		long long tmpK_ = K_;

		long long sub = static_cast<long long>(L[i]) * static_cast<long long>(C[i]);

		if (tmpK_ - sub <= 0) {
			ai = i;
			break;
		}
		else {
			K_ = tmpK_ - sub;
		}
	}

	// K_からLを引いていく
	long long K__ = K_;
	for (int i = 0; i < C[ai]; i++) {
		long long tmpK__ = K__;

		long long sub = static_cast<long long>(L[ai]);

		if (tmpK__ - sub <= 0) {
			break;
		}
		else {
			K__ = tmpK__ - sub;
		}
	}

	if (K__ == 0) {
		K__ = L[ai];
	}

	std::cout << A[ai][K__ - 1] << std::endl;
	return 0;
}

解説のコード

Ci LiのどのグループにKが存在するか見極めた後、%を使って、インデックスを求めている
Kは始まりにしていることに注意
基本特別な理由がない限り、0-indexで考えるのがよさそう

#include <iostream>
#include <string>
#include <map>
#include <unordered_map>
#include <cstring>
#include <cmath>
#include <algorithm>
#include <vector>
#include <queue>
#include <stack>
#include <limits.h>
#include <bitset>
#include <list>
#include <set>
#include <numeric>
#include <tuple>

int N; // 1~
std::vector<int> L; // Li : 整数列Aiの長さ
std::vector<std::vector<int>> A; // 整数列 各値は、1~10^9
long long K; // 1~ N 10^9 2x10^5
std::vector<int> C;
// Liの合計値は、2x10^5までなため、整数列A内のすべての要素を見ることはできる

// Bの末尾に、Aiを連結させる操作をCi回 i=1,2,...Nの順
// Bの長さは最大ΣCiLi
std::vector<int> B;

int main()
{
	std::cin.tie(0);
	std::ios::sync_with_stdio(false);

	std::cin >> N >> K;
	A.resize(N);
	L.resize(N);
	C.resize(N);

	for (int i = 0; i < N; i++) {
		std::cin >> L[i];
		A[i].resize(L[i]);
		for (int l = 0; l < L[i]; l++) {
			int a;
			std::cin >> a;
			A[i][l] = a;
		}
	}

	for (int i = 0; i < N; i++) {
		std::cin >> C[i];
	}

	K--; // Kを0-indexedにする!
	for (int i = 0; i < N; i++) {
		long long sub = static_cast<long long>(L[i]) * static_cast<long long>(C[i]);
		if (K < sub) {
			long long idx = K % static_cast<long long>(L[i]);
			std::cout << A[i][idx] << std::endl;
			return 0;
		}

		K -= sub;
	}
	return 0;
}