torus711 のアレ

主に競技プログラミングの問題について書きます.PC 以外だと数式が表示されないかもしれないです

TopCoder SRM 671, Division 1, Level 2 : BearDarts

問題概要

 正整数からなる列 $w$ が与えられる.4 要素からなる $w$ の部分列をとって $\{ a, b, c, d \}$ としたとき,$ac = bd$ となっているものの総数を求めよ.

  • $4 \leq |w| \leq 1{,}000$
  • $1 \leq w_i \leq 10^9$

解法

 数式が与えられているのでとりあえず変形してみます.すると $ac = bd \Leftrightarrow \frac a b = \frac d c$ であることが分かります.$a, b$ の組と $c, d$ の組で位置が交差しないので,こちらの方が扱いやすくなります.交差しない位置関係で考えると,$a, b$ をとる位置を決めたときに $b$ の位置より後ろから $\frac a b$ をいくつ作れるかを考える問題に帰着できます.$a, b$ の位置の取り方が $O( |w|^2 )$ 程度なので,それぞれについて $o( |w| )$ 時間で個数を求められれば Time Limit に間に合いそうです.

 個数の数え方ですが,まず単純な配列の場合について考えてみます.この場合,配列がソート済みならば,対象の値以上になる最初の位置と,超過になる最初の位置をそれぞれ二分探索で求めて距離をとると分かります.std::upper_bound と std::lower_boud のことです(std::equal_range でまとめてできる).
 元の問題に戻って考えると,ある $c$ の取り方に対して,それより右の全ての位置が $d$ の位置の候補なので,$c$ の位置に対して $\frac d c$ の形で書ける値の配列が 1 つ決まります.この配列をそれぞれソート済みにしておけば,$c$ の位置も決まった条件下であれば個数を数えられます.ただしこれだとまだ不十分で,更に工夫がいります.
 2 本のソート済み配列について考えてみると,マージソートの内部処理 (std::merge) のようにマージすると,2 本の配列を合併したソート済み配列が得られます.これに対して二分探索をすると,元の 2 つの配列に含まれる対象の個数が分かります.つまり,$c$ の取り方 2 通りを試して個数を数えるのと同じ結果を得られます.従って,幾つかのマージ済み配列を持っておいて,決め打ちした $b$ の位置より右を重複無くカバーできれば,$a, b$ を決め打ちしたときにとれる $\frac a b = \frac d c$ なる $c, d$ の個数を数えることができます.
 これは,ソート済み配列を Segment Tree の各ノードに乗せることで実現できます.親のノードでは,子ノードをマージした列を持っておきます.

 Segment Tree の元となる配列たちには $c, d$ の取り方と同じだけの個数の要素が含まれるので,構築に $O( |w|^2 )$ 時間かかります.Segment Tree の構築には,配列に出現する値が $O( |w|^2 )$ 個あって,構築時にそれぞれ $O( \log |w| )$ 回のマージに巻き込まれるので $O( |w|^2 \log |w| )$ 時間です.配列上での区間は Segment Tree の $O( \log |w| )$ 個のノードでカバーでき,その各ノードで二分探索を定数回行うので,$O( \log^2 |w| )$ 時間で連続する区間内に含まれるある値の個数を数えることができます.$a, b$ の取り方が $O( |w|^2 )$ 通りなので個数を数えるパートは $O( |w|^2 \log^2 |w| )$ 時間です.結局,$a, b$ の決め打ち + Segmment Tree へのクエリ部分に最も時間がかかって,全体では $O( |w|^2 \log^2 |w| )$ 時間となります.

 ところで,$\frac a b$ や $\frac d c$ の持ち方ですが,適切に扱わないと WA になります.有理数型を作るのが安全ですが,実数でも EPS を適切に設定して比較すれば(テストケースの範囲では)正しく数えられます.精度は 10 進で 18 桁ぐらい必要で,double 型だと精度が足りません.

コード

using LL = long long;
template < typename T = int > using VT = vector< T >;
template < typename T = int > using VVT = vector< vector< T > >;

#define REP2( i, n ) REP3( i, 0, n )
#define REP3( i, m, n ) for ( int i = ( int )( m ); i < ( int )( n ); ++i )
#define GET_REP( a, b, c, F, ... ) F
#define REP( ... ) GET_REP( __VA_ARGS__, REP3, REP2 )( __VA_ARGS__ )
#define ALL( c ) begin( c ), end( c )

#define SZ( v ) ( (int)( v ).size() )
#define EB emplace_back
#define BI back_inserter

template < typename T = long long >
class Rational
{
	T P, Q;

public:
	Rational( const T p, const T q ) :
		Rational( p, q, __gcd( p, q ) )
	{
		return;
	}

	Rational( const T p, const T q, const T g ) :
		P( p / g ), Q( q / g )
	{
		return;
	}

	bool operator<( const Rational &rhs ) const
	{
		return P * rhs.Q < rhs.P * Q;
	}
};

// セグメント木のユーティリティ
inline int chl( const int k )
{
	return k * 2 + 1;
}

inline int chr( const int k )
{
	return k * 2 + 2;
}

inline int mid( const int l, const int r )
{
	return ( l + r ) / 2;
}

#define LEFT( k, l, r ) chl( k ), l, mid( l, r )
#define RIGHT( k, l, r ) chr( k ), mid( l, r ), r

struct SegmentTree
{
	const int N;
	VVT< Rational<> > data;

	SegmentTree( const VVT< Rational<> > &src ) :
		N( SZ( src ) ), data( N * 4 )
	{
		init( src, 0, 0, N );
		return;
	}

	void init( const VVT< Rational<> > &src, const int k, const int l, const int r )
	{
		if ( l + 1 == r )
		{
			data[k] = src[l];
			return;
		}

		init( src, LEFT( k, l, r ) );
		init( src, RIGHT( k, l, r ) );

		merge( ALL( data[ chl( k ) ] ), ALL( data[ chr( k ) ] ), BI( data[k] ) );

		return;
	}

	int query( const int a, const int b, const Rational<> &t )
	{
		return query( a, b, t, 0, 0, N );
	}
	int query( const int a, const int b, const Rational<> &t, const int k, const int l, const int r )
	{
		if ( r <= a || b <= l )
		{
			return 0;
		}
		else if ( a <= l && r <= b )
		{
			const auto lb = lower_bound( ALL( data[k] ), t );
			const auto ub = upper_bound( ALL( data[k] ), t );
			return ub - lb;
		}
		else
		{
			return query( a, b, t, LEFT( k, l, r ) ) + query( a, b, t, RIGHT( k, l, r ) );
		}
	}
};


class BearDarts
{
public:
	long long count( vector <int> W )
	{
		const int N = SZ( W );

		VVT< Rational<> > arrays( N );
		REP( i, N )
		{
			REP( j, i + 1, N )
			{
				arrays[i].EB( W[j], W[i] );
			}
		}
		for_each( ALL( arrays ), []( VT< Rational<> > &a ){ sort( ALL( a ) ); } );

		SegmentTree segtree( move( arrays ) );

		LL res = 0;
		REP( i, N )
		{
			REP( j, i + 1, N )
			{
				res += segtree.query( j + 1, N, Rational<>( W[i], W[j] ) );
			}
		}

		return res;
	}
};