Codechef - Colored Array

color.cpp

#include <cstdio>
#include <cstdlib>
#include <cstring>
#include <algorithm>
#include <functional>

#define MAXN 1001

using namespace std;

int t, k, n, m, total;
int A[MAXN], B[MAXN][MAXN], C[MAXN][MAXN], haruhi[MAXN];

int main() {
	scanf("%d", &t);
while(t--) {
	total = 0;

	scanf("%d %d %d", &n, &m, &k);
	for(int i = 1; i <= n; i++) scanf("%d", &A[i]);
	for(int i = 1; i <= n; i++)
		for(int j = 1; j <= m; j++)
			scanf("%d", &B[i][j]);
	for(int i = 1; i <= n; i++)
		for(int j = 1; j <= m; j++)
			scanf("%d", &C[i][j]);

	for(int i = 1; i <= n; i++)
		total += B[i][A[i]];

	for(int i = 0; i <= n; i++)
		haruhi[i] = 0;
	for(int i = 1; i <= n; i++)
		for(int j = 1; j <= m; j++)
			if(j != A[i])
				haruhi[i] = max(haruhi[i], B[i][j]-B[i][A[i]] - C[i][j]);

	sort(haruhi+1,haruhi+n+1,greater<int>());

	for(int i = 1; i <= k && i <= n && haruhi[i] > 0; i++)
		total += haruhi[i];

	printf("%d\n", total);
}
	return 0;
}