Submission #736154

#	Time	Username	Problem	Language	Result	Execution time	Memory
736154		marvinthang	Janjetina (COCI21_janjetina)	C++17	110 / 110	425 ms	16572 KiB

Main

/*************************************
*    author: marvinthang             *
*    created: 05.08.2022 16:14:49    *
*************************************/

#include <bits/stdc++.h>

using namespace std;

#define                  fi  first
#define                  se  second
#define                 div  ___div
#define                left  ___left
#define               right  ___right
#define                TIME  (1.0 * clock() / CLOCKS_PER_SEC)
#define             MASK(i)  (1LL << (i))
#define             FULL(i)  (MASK(i) - 1)
#define           BIT(x, i)  ((x) >> (i) & 1)
#define  __builtin_popcount  __builtin_popcountll
#define        scan_op(...)  istream & operator >> (istream &in, __VA_ARGS__ &u)
#define       print_op(...)  ostream & operator << (ostream &out, const __VA_ARGS__ &u)
#define          file(name)  if (fopen (name".inp", "r")) { freopen (name".inp", "r", stdin); freopen (name".out", "w", stdout); }

template <class T>             scan_op(vector <T>)  { for (size_t i = 0; i < u.size(); ++i) in >> u[i]; return in; }
template <class T>            print_op(vector <T>)  { out << '{'; for (size_t i = 0; i + 1 < u.size(); ++i) out << u[i] << ", "; if (!u.empty()) out << u.back(); return out << '}'; }
template <class U, class V>   scan_op(pair <U, V>)  { return in >> u.fi >> u.se; }
template <class U, class V>  print_op(pair <U, V>)  { return out << '(' << u.fi << ", " << u.se << ')'; }

const double PI = 3.1415926535897932384626433832795; // acos(-1.0); atan(-1.0);
const int dir[] = {1, 0, -1, 0, 1, 1, -1, -1, 1}; // {2, 1, -2, -1, -2, 1, 2, -1, 2};
const int MAX = 1e5 + 5;

// index from 0
template <class T>
    struct Fenwick {

        T *bit;
        int n = -1;

        Fenwick(int _n = 0) {
            resize(_n);
        }

        void reset(void) {
            fill(bit, bit + 1 + n, T(0));
        }

        void resize(int _n) {
            if (~n) delete bit;
            n = _n;
            bit = new T[n + 1];
            reset();
        }

        void update(int i, T val) {
            assert(0 <= i);
            ++i;
            for (; i <= n; i += i & -i) bit[i] += val;
        }

        void update(int l, int r, T val) {
            if (l > r) return;
            update(l, val);
            update(r + 1, -val);
        }

        T get(int i) {
            if (i < 0) return T(0);
            ++i;
            i = min(i, n);
            T res(0);
            for (; i > 0; i -= i & -i)
                res += bit[i];
            return res;
        }

        T get(int l, int r) {
            if (l > r) return T(0);
            return get(r) - get(l - 1);
        }

        int upper_bound(T val) {
            int res = 0;
            for (int i = __lg(n); i >= 0; --i) {
                if ((res | MASK(i)) <= n && val >= bit[res | MASK(i)]) {
                    res |= MASK(i);
                    val -= bit[res];
                }
            }
            return res;
        }

        int lower_bound(T val) {
            int res = 0;
            for (int i = __lg(n); i >= 0; --i) {
                if ((res | MASK(i)) <= n && val > bit[res | MASK(i)]) {
                    res |= MASK(i);
                    val -= bit[res];
                }
            }
            return res;
        }

    };


int N, K;
vector <pair <int, int>> adj[MAX];
Fenwick <int> bit;

void loadTree(void) {
	cin >> N >> K;
	for (int i = 1; i < N; ++i) {
		int u, v, w; cin >> u >> v >> w;
		adj[u].push_back(make_pair(w, v));
		adj[v].push_back(make_pair(w, u));
	}
}

bool used[MAX];
int treeSize[MAX];

int calcSize(int u, int par) {
	treeSize[u] = 1;
	for (auto [_, v]: adj[u]) if (!used[v] && v != par) treeSize[u] += calcSize(v, u);
	return treeSize[u];
}

int findCentroid(int u, int par, int n) {
	for (auto [_, v]: adj[u]) if (!used[v] && v != par && treeSize[v] > n / 2) return findCentroid(v, u, n);
	return u;
}


void dfs(int u, int par, int dep, int ma, vector <pair <int, int>> &values) {
	values.push_back(make_pair(ma, dep));
	for (auto [w, v]: adj[u]) if (!used[v] && v != par) 
		dfs(v, u, dep + 1, max(ma, w), values);
}

long long solve(int u) {
	int n = calcSize(u, 0);
	u = findCentroid(u, 0, n);

	vector <pair <int, int>> values;
	dfs(u, 0, 0, 0, values);
	sort(values.begin(), values.end());
	long long res = 0;
	for (auto [ma, dep]: values) {
		res += bit.get(ma - dep - K);
		bit.update(dep, 1);
	}
	for (auto [ma, dep]: values) bit.update(dep, -1);
	
	for (auto [w, v]: adj[u]) if (!used[v]) {
		values.clear();
		dfs(v, u, 1, w, values);
		sort(values.begin(), values.end());
		for (auto [ma, dep]: values) {
			res -= bit.get(ma - dep - K);
			bit.update(dep, 1);
		}
		for (auto [ma, dep]: values) bit.update(dep, -1);
	}

	used[u] = true;
	for (auto [_, v]: adj[u]) if (!used[v]) res += solve(v);
	return res;
}

void process(void) {
	bit.resize(MAX);
	cout << solve(1) * 2 << '\n';
}

int main(void) {
	ios_base::sync_with_stdio(false); cin.tie(nullptr); // cout.tie(nullptr);
	file("coci2021_r4_janjetina");
	loadTree();
	process();
	// cerr << "Time elapsed: " << TIME << " s.\n";
	return (0^0);
}

Compilation message (stderr)

Main.cpp: In function 'int main()':
Main.cpp:22:69: warning: ignoring return value of 'FILE* freopen(const char*, const char*, FILE*)' declared with attribute 'warn_unused_result' [-Wunused-result]
   22 | #define          file(name)  if (fopen (name".inp", "r")) { freopen (name".inp", "r", stdin); freopen (name".out", "w", stdout); }
      |                                                             ~~~~~~~~^~~~~~~~~~~~~~~~~~~~~~~~
Main.cpp:178:2: note: in expansion of macro 'file'
  178 |  file("coci2021_r4_janjetina");
      |  ^~~~
Main.cpp:22:103: warning: ignoring return value of 'FILE* freopen(const char*, const char*, FILE*)' declared with attribute 'warn_unused_result' [-Wunused-result]
   22 | #define          file(name)  if (fopen (name".inp", "r")) { freopen (name".inp", "r", stdin); freopen (name".out", "w", stdout); }
      |                                                                                               ~~~~~~~~^~~~~~~~~~~~~~~~~~~~~~~~~
Main.cpp:178:2: note: in expansion of macro 'file'
  178 |  file("coci2021_r4_janjetina");
      |  ^~~~

• Subtask 1: 15 / 15
• Subtask 2: 35 / 35
• Subtask 3: 60 / 60