""" 条件概率计算工具： 1）计算某个人设树节点在父节点下的条件概率； 2）计算某个 pattern 的条件概率。 """ from __future__ import annotations import itertools import json from pathlib import Path from typing import Any # 节点名 -> (该节点 post_ids, 父节点 post_ids)，用 frozenset 便于批量计算时复用、避免重复转换 NodePostIndex = dict[str, tuple[frozenset[str], frozenset[str]]] # 已推导列表：每项为 (已推导的选题点, 推导来源人设树节点)，如 ("分享","分享")、("柴犬","动物角色") # 推导来源人设树节点的 post_ids 在计算条件概率时从人设树中读取 DerivedItem = tuple[str, str] def _tree_dir(account_name: str, base_dir: Path | None = None) -> Path: """人设树目录：../input/{account_name}/处理后数据/tree/（相对本文件所在目录）。""" if base_dir is not None: return base_dir / account_name / "处理后数据" / "tree" return Path(__file__).resolve().parent.parent / "input" / account_name / "处理后数据" / "tree" def _load_trees(account_name: str, base_dir: Path | None = None) -> list[tuple[str, dict]]: """加载该账号下所有维度的人设树。返回 [(维度名, 根节点 dict), ...]。""" td = _tree_dir(account_name, base_dir) if not td.is_dir(): return [] return _load_trees_from_directory(td) def _load_trees_from_directory(tree_dir: Path) -> list[tuple[str, dict]]: """ 从指定目录加载所有人设树 JSON（每文件取顶层第一个维度根，与按账号目录加载时行为一致）。 用于平台库等人设树路径非 input/{账号}/处理后数据/tree/ 的场景。 """ if not tree_dir.is_dir(): return [] result: list[tuple[str, dict]] = [] for p in sorted(tree_dir.glob("*.json")): try: with open(p, "r", encoding="utf-8") as f: data = json.load(f) for dim_name, root in data.items(): if isinstance(root, dict): result.append((str(dim_name), root)) break except Exception: continue return result def _post_ids_of(node: dict) -> list[str]: """从树节点中取出 _post_ids，无则返回空列表。""" return list(node.get("_post_ids") or []) def _build_node_index_from_trees(trees: list[tuple[str, dict]]) -> dict[str, tuple[list[str], list[str]]]: """ 遍历多棵人设树，建立 节点名 -> (该节点 post_ids, 父节点 post_ids)。 同一节点名在多个分支出现时，保留第一次遇到的（保证父子一致）。 """ index: dict[str, tuple[list[str], list[str]]] = {} for _dim, root in trees: parent_pids = _post_ids_of(root) def walk(parent_ids: list[str], node_dict: dict) -> None: for name, child in (node_dict.get("children") or {}).items(): if not isinstance(child, dict): continue if name not in index: index[name] = (_post_ids_of(child), list(parent_ids)) walk(_post_ids_of(child), child) walk(parent_pids, root) return index def _build_node_index(account_name: str, base_dir: Path | None = None) -> dict[str, tuple[list[str], list[str]]]: """遍历账号下所有人设树，建立节点索引。""" return _build_node_index_from_trees(_load_trees(account_name, base_dir)) def build_node_post_index_from_tree_dir(tree_dir: Path) -> NodePostIndex: """从任意人设树目录（如 input/xiaohongshu/tree）构建节点 post 索引，算法与账号树一致。""" raw = _build_node_index_from_trees(_load_trees_from_directory(tree_dir)) return {k: (frozenset(a), frozenset(b)) for k, (a, b) in raw.items()} def build_node_index_for_tree_dir(tree_dir: Path) -> dict[str, tuple[list[str], list[str]]]: """从任意人设树目录构建节点名 -> (节点 post_ids, 父 post_ids)，供 pattern 条件概率等使用。""" return _build_node_index_from_trees(_load_trees_from_directory(tree_dir)) def load_persona_trees_from_dir(tree_dir: Path) -> list[tuple[str, dict]]: """从目录加载人设树列表（每 JSON 文件一个顶层维度），供遍历节点等场景复用。""" return _load_trees_from_directory(tree_dir) def _derived_post_ids_from_sources( derived_list: list[DerivedItem], index: dict[str, tuple[list[str], list[str]]], ) -> set[str]: """根据 derived_list 中的「推导来源人设树节点」在人设树中的 post_ids 取交集，得到已推导的帖子集合。""" common: set[str] | None = None for _topic_point, source_node in derived_list: if source_node not in index: continue pids = set(index[source_node][0]) if common is None: common = pids else: common &= pids return common if common is not None else set() def _derived_post_ids_from_frozen_index( derived_list: list[DerivedItem], index: NodePostIndex, ) -> frozenset[str]: """与 _derived_post_ids_from_sources 相同语义，索引为 frozenset 版（批量场景复用）。""" common: frozenset[str] | None = None for _topic_point, source_node in derived_list: if source_node not in index: continue pids = index[source_node][0] common = pids if common is None else common & pids return common if common is not None else frozenset() def build_node_post_index(account_name: str, base_dir: Path | None = None) -> NodePostIndex: """ 构建账号人设树的节点索引（每个节点只建一次，供批量 calc_node_conditional_ratio 复用）。 值为 (节点 post_ids, 父节点 post_ids) 的 frozenset，减少重复 list->set 与拷贝。 """ raw = _build_node_index(account_name, base_dir) return {k: (frozenset(a), frozenset(b)) for k, (a, b) in raw.items()} def calc_node_conditional_ratio( account_name: str, derived_list: list[DerivedItem], tree_node_name: str, base_dir: Path | None = None, node_post_index: NodePostIndex | None = None, target_ratio: float | None = None, ) -> float: """ 计算人设树节点 N 在父节点 P 下的条件概率。 参数: account_name: 账号名称 derived_list: 已推导列表，每项 (已推导的选题点, 推导来源人设树节点) tree_node_name: 人设树节点 N 的名称（字符串匹配） base_dir: 可选，input 根目录；不传则使用相对本文件的 ../input node_post_index: 可选，由 build_node_post_index 预构建；批量对多节点计算时传入可避免重复读盘与遍历整棵树 target_ratio: 可选，目标条件概率。若某个组合的条件概率已达到该值，则直接返回（用于缩小组合搜索） 计算规则: 已推导的帖子集合：从 derived_list 中先取「最多选题点」的交集，再逐步减少到 1 个选题点， 对每种选题点子集分别计算条件概率，最后取最大值。 对每种情况：已推导的帖子集合 = 该子集中各「推导来源人设树节点」在人设树中的 post_ids 的交集； 分子 = |已推导的帖子集合 ∩ N 的 post_ids|，分母 = |已推导的帖子集合 ∩ P 的 post_ids|； 条件概率 = 分子/分母，且 ≤1；分母为 0 时该情况跳过。 """ index = node_post_index if node_post_index is not None else build_node_post_index(account_name, base_dir) if tree_node_name not in index: return 0.0 set_n, set_p = index[tree_node_name] # 关键优化（不改变搜索空间/结果）： # - derived_list 里重复的 source_node 对“交集”没有任何影响，但会把 L 变大导致 2^L 爆炸 # - 不在 index 里的 source_node 原本也会被跳过，提前过滤可减少组合规模 # - 组合内交集直接对 frozenset 逐步 &，避免 list(combo)/函数调用开销 seen_sources: set[str] = set() source_sets: list[frozenset[str]] = [] for _topic, source_node in derived_list: if source_node in seen_sources: continue seen_sources.add(source_node) tup = index.get(source_node) if tup is None: continue source_sets.append(tup[0]) if not source_sets: return 0.0 # 将更小的集合放前面：交集会更快“变小”，每次 & 的成本更低（仍然枚举全部子集） source_sets.sort(key=len) max_ratio = 0.0 # 从 1 个选题点到「最多选题点」：对每种子集大小，取所有组合，分别算条件概率后取最大 for k in range(1, len(source_sets) + 1): for combo_sets in itertools.combinations(source_sets, k): derived_post_ids = combo_sets[0] for s in combo_sets[1:]: derived_post_ids = derived_post_ids & s den = len(derived_post_ids & set_p) if den == 0: continue num = len(derived_post_ids & set_n) ratio = min(1.0, num / den) max_ratio = max(max_ratio, ratio) if target_ratio is not None and max_ratio >= target_ratio: return round(max_ratio, 4) return round(max_ratio, 4) def _pattern_nodes_and_post_count(pattern: dict[str, Any]) -> tuple[list[str], int, float]: """ 从 pattern 中解析出节点列表和 post_count。支持 nodes + post_count 或 i + post_count。 返回的 post_count 表示该 pattern 本身的帖子数，在条件概率计算中作为分子（即 pattern 本身的概率/占比的分子）。 """ nodes = pattern.get("nodes") if nodes is not None and isinstance(nodes, list): nodes = [str(x).strip() for x in nodes if x] else: raw = pattern.get("i") or pattern.get("pattern_str") or "" nodes = [x.strip() for x in str(raw).replace("+", " ").split() if x.strip()] post_count = int(pattern.get("post_count", 0)) support = pattern.get("s", 0.0) return nodes, post_count, support def calc_pattern_conditional_ratio_with_index( derived_list: list[DerivedItem], pattern: dict[str, Any], index: dict[str, tuple[list[str], list[str]]], ) -> float: """ 与 calc_pattern_conditional_ratio 相同计算规则，但使用已构建的人设树节点索引 （例如平台库 input/xiaohongshu/tree）。 """ pattern_nodes, post_count, pattern_s = _pattern_nodes_and_post_count(pattern) if not pattern_nodes or post_count <= 0: return pattern_s derived_sources = set(source for _post, source in derived_list) derived_pattern_nodes = [n for n in pattern_nodes if n in derived_sources] if not derived_pattern_nodes: return pattern_s derived_in_tree = [n for n in derived_pattern_nodes if n in index] if not derived_in_tree: return pattern_s common: set[str] | None = None for name in derived_in_tree: pids = set(index[name][0]) if common is None: common = pids else: common &= pids if common is None or len(common) == 0: return pattern_s den = len(common) return round(min(1.0, post_count / den), 4) def calc_pattern_conditional_ratio( account_name: str, derived_list: list[DerivedItem], pattern: dict[str, Any], base_dir: Path | None = None, ) -> float: """ 计算某个 pattern 的条件概率。 参数: account_name: 账号名称 derived_list: 已推导列表，每项 (已推导的选题点, 推导来源人设树节点) pattern: 至少包含节点列表与 post_count。 - 节点列表: key 为 "nodes"（list）或 "i"（字符串，用 + 连接） - post_count: 该 pattern 的帖子数量，作为分子 base_dir: 可选，input 根目录 计算规则: 取 pattern 中「已被推导」的节点（其名称出现在 derived 的推导来源中）， 在人设树中取这些节点的 post_ids 的交集作为分母； 分子 = pattern.post_count（由 _pattern_nodes_and_post_count 解析得到，表示 pattern 本身的帖子数）。 条件概率 = 分子/分母，且 ≤1；分母为 0 时返回 1。 """ index = _build_node_index(account_name, base_dir) return calc_pattern_conditional_ratio_with_index(derived_list, pattern, index) def _test_with_user_example() -> None: """ 使用你提供的测试数据：已推导 (分享|分享)、(柴犬|动物角色)； 人设树节点：恶作剧；pattern：分享+动物角色+创意表达 post_count=2。 推导来源的 post_ids 在方法内部从人设树读取。 """ account_name = "阿里多多酱" # 已推导列表：(已推导的选题点, 推导来源人设树节点) derived_list: list[DerivedItem] = [ ("推广", "推广"), ("视觉调性", "视觉调性"), # ("图片文字", "图片文字"), # ("补充说明式", "补充说明式"), # ("幽默化标题", "幽默化标题"), # ("标题", "标题"), ] # 1）人设树节点「恶作剧」的条件概率 r_node = calc_node_conditional_ratio(account_name, derived_list, "观念") print(f"1) 人设树节点条件概率: {r_node}") # 2）pattern 分享+动物角色+创意表达 post_count=2 的条件概率 pattern = {"i": "视觉调性+辞格意象+叙事编排", "post_count": 22, "s": 0.478261} r_pattern = calc_pattern_conditional_ratio(account_name, derived_list, pattern) print(f"2) pattern 条件概率: {r_pattern}") if __name__ == "__main__": _test_with_user_example()