| 258 | |
| 259 | |
| 260 | TVector<TArraySubsetIndexing<ui32>> NCB::Split( |
| 261 | const TObjectsGrouping& objectsGrouping, |
| 262 | ui32 partCount, |
| 263 | bool oldCvStyle |
| 264 | ) { |
| 265 | const ui32 objectCount = objectsGrouping.GetObjectCount(); |
| 266 | |
| 267 | TVector<TArraySubsetIndexing<ui32>> result; |
| 268 | |
| 269 | if (objectsGrouping.IsTrivial()) { |
| 270 | ui32 currentEnd = 0; |
| 271 | for (ui32 part = 0; part < partCount; ++part) { |
| 272 | TSubsetBlock<ui32> block; |
| 273 | if (oldCvStyle) { |
| 274 | block.SrcBegin = currentEnd; |
| 275 | block.SrcEnd = objectCount * (part + 1) / partCount; |
| 276 | currentEnd = block.SrcEnd; |
| 277 | } else { |
| 278 | InitElementRange( |
| 279 | part, |
| 280 | partCount, |
| 281 | objectCount, |
| 282 | &block.SrcBegin, |
| 283 | &block.SrcEnd |
| 284 | ); |
| 285 | } |
| 286 | const ui32 blockSize = block.GetSize(); |
| 287 | CB_ENSURE(blockSize > 0, "Not enough objects for splitting into requested amount of parts"); |
| 288 | block.DstBegin = 0; |
| 289 | result.push_back( |
| 290 | TArraySubsetIndexing<ui32>( |
| 291 | TRangesSubset<ui32>(blockSize, TVector<TSubsetBlock<ui32>>{std::move(block)}) |
| 292 | ) |
| 293 | ); |
| 294 | } |
| 295 | } else { |
| 296 | const ui32 partSize = objectsGrouping.GetObjectCount() / partCount; |
| 297 | ui32 currentPartObjectEnd = 0; |
| 298 | ui32 currentPartGroupEnd = 0; |
| 299 | for (ui32 part = 0; part < partCount; ++part) { |
| 300 | currentPartObjectEnd = |
| 301 | oldCvStyle ? |
| 302 | (objectCount * (part + 1) / partCount) : |
| 303 | Min(currentPartObjectEnd + partSize, objectCount); |
| 304 | const ui32 lastGroupIdx = ( |
| 305 | part + 1 == partCount ? |
| 306 | (objectsGrouping.GetGroupCount() - 1) : |
| 307 | objectsGrouping.GetGroupIdxForObject(currentPartObjectEnd - 1) |
| 308 | ); |
| 309 | TSubsetBlock<ui32> block{{currentPartGroupEnd, lastGroupIdx + 1}, 0}; |
| 310 | const ui32 blockSize = block.GetSize(); |
| 311 | CB_ENSURE(blockSize > 0, "Not enough objects for splitting into requested amount of parts"); |
| 312 | result.push_back( |
| 313 | TArraySubsetIndexing<ui32>( |
| 314 | TRangesSubset<ui32>(blockSize, TVector<TSubsetBlock<ui32>>{std::move(block)}) |
| 315 | ) |
| 316 | ); |
| 317 | currentPartGroupEnd = lastGroupIdx + 1; |