結果

問題 No.364 門松木
コンテスト
ユーザー koyumeishi
提出日時 2016-03-24 00:23:09
言語 C++11（廃止可能性あり）
(gcc 13.3.0)
結果
AC  
実行時間 34 ms / 3,000 ms
コード長 2,624 bytes
コンパイル時間 1,126 ms
コンパイル使用メモリ 105,504 KB
実行使用メモリ 17,152 KB
最終ジャッジ日時 2024-10-01 21:41:42
合計ジャッジ時間 2,946 ms
ジャッジサーバーID
（参考情報） 		judge4 / judge2
このコードへのチャレンジ
（要ログイン）
ファイルパターン 結果
sample AC * 3
other AC * 31
権限があれば一括ダウンロードができます
コンパイルメッセージ
main.cpp: In function ‘std::istream& operator>>(std::istream&, std::vector<int>&)’:
main.cpp:18:78: warning: ignoring return value of ‘int scanf(const char*, ...)’ declared with attribute ‘warn_unused_result’ [-Wunused-result]
   18 | istream& operator >> (istream& is, vector<int>& vec){for(int& val: vec) scanf("%d", &val); return is;}
      |                                                                         ~~~~~^~~~~~~~~~~~
main.cpp: In function ‘std::istream& operator>>(std::istream&, std::vector<long long int>&)’:
main.cpp:19:90: warning: ignoring return value of ‘int scanf(const char*, ...)’ declared with attribute ‘warn_unused_result’ [-Wunused-result]
   19 | istream& operator >> (istream& is, vector<long long>& vec){for(long long& val: vec) scanf("%lld", &val); return is;}
      |                                                                                     ~~~~~^~~~~~~~~~~~~~
main.cpp: In function ‘std::istream& operator>>(std::istream&, std::vector<double>&)’:
main.cpp:20:84: warning: ignoring return value of ‘int scanf(const char*, ...)’ declared with attribute ‘warn_unused_result’ [-Wunused-result]
   20 | istream& operator >> (istream& is, vector<double>& vec){for(double& val: vec) scanf("%lf", &val); return is;}
      |                                                                               ~~~~~^~~~~~~~~~~~~
main.cpp: In function ‘std::istream& operator>>(std::istream&, std::vector<std::__cxx11::basic_string<char> >&)’:
main.cpp:22:85: warning: ignoring return value of ‘int scanf(const char*, ...)’ declared with attribute ‘warn_unused_result’ [-Wunused-result]
   22 | istream& operator >> (istream& is, vector<string>& vec){for(string& val: vec) {scanf("%s", buff); val=buff;}; return is;}
      |                                                                                ~~~~~^~~~~~~~~~~~
main.cpp: In function ‘int main()’:
main.cpp:87:22: warning: ignoring return value of ‘int

ソースコード

diff # 

#include <iostream>
#include <vector>
#include <cstdio>
#include <sstream>
#include <map>
#include <string>
#include <algorithm>
#include <queue>
#include <cmath>
#include <functional>
#include <set>
#include <ctime>
#include <random>
#include <chrono>
using namespace std;

//input from stdin
istream& operator >> (istream& is, vector<int>& vec){for(int& val: vec) scanf("%d", &val); return is;}
istream& operator >> (istream& is, vector<long long>& vec){for(long long& val: vec) scanf("%lld", &val); return is;}
istream& operator >> (istream& is, vector<double>& vec){for(double& val: vec) scanf("%lf", &val); return is;}
char buff[200000];
istream& operator >> (istream& is, vector<string>& vec){for(string& val: vec) {scanf("%s", buff); val=buff;}; return is;}

template<class T> istream& operator >> (istream& is, vector<T>& vec){for(T& val: vec) is >> val; return is;}
template<class T> istream& operator ,  (istream& is, T& val){ return is >> val;}
template<class T> ostream& operator << (ostream& os, const vector<T>& vec){for(int i=0; i<vec.size(); i++) os << vec[i] << (i==vec.size()-1?"":" "); return os;}
template<class T> ostream& operator ,  (ostream& os, const T& val){ return os << " " << val;}
template<class T> ostream& operator >> (ostream& os, const T& val){ return os << " " << val;}

long long dfs(vector<int>& a, vector<vector<int>>& G, int pos, int last, long long& ans){
	map<int, long long> dp;

	for(int next : G[pos]){
		if(next == last) continue;
		long long tmp = dfs(a,G, next, pos, ans);
		if(a[next] == a[pos]) continue;

		if(dp.find(a[next]) != dp.end()) dp[a[next]] = max(dp[a[next]], tmp);
		else dp[a[next]] = tmp;
		ans = max(ans, tmp);
	}

	long long l = 0;
	long long l_cnt = 0;
	long long s = 0;
	long long s_cnt = 0;
	for(auto v : dp){
		if(v.first < a[pos]){
			s += v.second;
			s_cnt++;
		}else if(v.first > a[pos]){
			l += v.second;
			l_cnt++;
		}
	}

	ans = max({ans, l + l_cnt*(l_cnt-1)/2, s + s_cnt*(s_cnt-1)/2});

	if(last == pos) return 0;
	if(a[pos] == a[last]) return 0;
	if(a[pos] < a[last]){
		if(dp.find(a[last]) != dp.end()){
			l -= dp[a[last]];
			l_cnt--;
		}
		return l + l_cnt*(l_cnt-1)/2 + l_cnt;
	}else if(a[pos] > a[last]){
		if(dp.find(a[last]) != dp.end()){
			s -= dp[a[last]];
			s_cnt--;
		}
		return s + s_cnt*(s_cnt-1)/2 + s_cnt;
	}
	return 0;
}


int main(){
	int n;
	cin >> n;
	vector<int> a(n);
	cin >> a;
	vector<vector<int>> G(n);
	for(int i=0; i<n-1; i++){
		int u,v;
		//cin >> u,v;
		scanf("%d%d", &u,&v);
		u--; v--;
		G[u].push_back(v);
		G[v].push_back(u);
	}

	long long ans = 0;
	dfs(a,G, 0,0, ans);

	cout << ans << endl;
	return 0;
}
0