torus711 のアレ

主に競技プログラミングの問題について書きます

Codeforces 167, Division 2, B : Dima and Sequence

概要

n 項からなる数列と、次のように定義される関数 f が与えられる。

  • f(0) = 0
  • f(2x) = f(x)
  • f(2x+1) = f(x) + 1

このとき、i < j かつ f( a_i ) = f( a_j ) を満たすような i, j の組の数を求めよ。

解法

まず、各 a_i について、f(a_i)の値を求めます。
関数の再帰で渡される引数は一回で約半分になるので、この処理は全体で O( n \quad \log \quad max(a) ) で終わります。
次に、結果が等しくなる要素の数を数え、それぞれについて組み合わせの数を求めて和を計算することで答えが求まります。

コード

typedef long long LL;
typedef vector<int> VI;

#define REP( i, m, n ) for ( int i = (int)( m ); i < (int)( n ); ++i )
#define EACH( v, c ) for ( auto &v : c )
#define EXIST( c, e ) ( (c).find( e ) != (c).end() )
#define snd second

int rec( int n )
{
	if ( n <= 0 )
	{
		return 0;
	}

	const int x = n / 2;

	if ( 2 * x == n )
	{
		return rec( x );
	}
	else
	{
		return rec( x ) + 1;
	}
}

LL sum( LL n )
{
	LL res = 0;

	REP( i, 1, n + 1 )
	{
		res += i;
	}

	return res;
}

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

	int n;
	cin >> n;

	VI as( n );
	EACH( a, as )
	{
		cin >> a;
	}

	map<int,int> f, count;

	REP( i, 0, n )
	{
		if ( !EXIST( f, as[i] ) )
		{
			f[ as[i] ] = rec( as[i] );
		}
		count[ f[ as[i] ] ] ++;
	}

	LL res = 0;

	EACH( c, count )
	{
		res += sum( c.snd - 1 );
	}

	cout << res << endl;

	return 0;
}